import mindspore as ms
from mindspore import nn, ops
import numpy as np

class CustomActivation(nn.Cell):
    def __init__(self):
        super().__init__()
        
    def construct(self, x):
        # 自定义激活函数：x * sin(x)
        return x * ops.sin(x)
    
    def bprop(self, x, out, dout):
        # 手动实现的反向传播：dout * (sin(x) + x * cos(x))
        return (dout * (ops.sin(x) + x * ops.cos(x)), )

# 测试代码
custom_op = CustomActivation()
x = ms.Tensor(np.random.randn(10), dtype=ms.float32)

# 自动梯度计算
grad_auto = ms.grad(custom_op)(x)

# 数值梯度验证
def numerical_gradient(f, x, eps=1e-6):
    grad = np.zeros_like(x.asnumpy())
    x_np = x.asnumpy()
    for i in range(x_np.size):
        x_plus = x_np.copy()
        x_minus = x_np.copy()
        x_plus[i] += eps
        x_minus[i] -= eps
        grad[i] = (f(ms.Tensor(x_plus)).asnumpy() - 
                   f(ms.Tensor(x_minus)).asnumpy()) / (2 * eps)
    return grad

def forward_func(x):
    return custom_op(x).sum()

grad_numerical = numerical_gradient(forward_func, x)
print(f"自动梯度: {grad_auto}")
print(f"数值梯度: {grad_numerical}")
print(f"最大差异: {np.abs(grad_auto.asnumpy() - grad_numerical).max()}")