Function tensor_grad_scale

codegeex/mindspore/src/pangu_alpha_wrapcell.py:72–73

(scale, grad)

@grad_scale.register("Tensor", "Tensor")
def tensor_grad_scale(scale, grad):
    return grad * P.Cast()(reciprocal(scale), F.dtype(grad))


@grad_scale.register("Tensor", "Tensor", "Tensor")