计算JAX中神经网络模型中梯度的点积自身的问题

使用 JAX 的 grad 和 vmap 函数以及 numpy 的 dot 函数实现此问题

import jax.numpy as np
from jax import grad, vmap

# 定义模型和输入
def model(params, x):
    return np.dot(params, x)

params = np.array([1.0, 2.0, 3.0])
inputs = np.array([[1.0, 2.0], [3.0, 4.0]])

# 定义函数计算梯度点积自身
def grad_dot_product(params, inputs):
    # 定义对模型的偏导数函数
    grad_fn = grad(model)
    # 对每个输入计算梯度
    grads = vmap(grad_fn, in_axes=(None, 0))(params, inputs)
    # 将每个梯度向量的点积自身相加
    return np.sum(np.dot(grads, grads.T))

# 调用函数计算结果
result = grad_dot_product(params, inputs)
print(result)