blas.py 文件源码

def backward(ctx, grad_output):
        vector1, vector2 = ctx.saved_variables
        grad_add_matrix = grad_vector1 = grad_vector2 = None

        if ctx.needs_input_grad[0]:
            grad_add_matrix = grad_output
            if ctx.alpha != 1:
                grad_add_matrix = grad_add_matrix.mul(ctx.alpha)

        if ctx.needs_input_grad[1]:
            grad_vector1 = torch.mv(grad_output, vector2)
            if ctx.beta != 1:
                grad_vector1 *= ctx.beta

        if ctx.needs_input_grad[2]:
            # TODO: maybe it's better to do transpose + mv + transpose
            grad_vector2 = torch.mm(vector1.unsqueeze(0), grad_output).squeeze(0)
            if ctx.beta != 1:
                grad_vector2 *= ctx.beta

        return grad_add_matrix, grad_vector1, grad_vector2, None, None, None