JAX没有提供任何方法来获取单个矩阵元素的梯度。有两种方法可以继续;第一，你可以获取整个数组的梯度并提取你感兴趣的元素;例如：

import jax
import jax.numpy as jnp

def f(A):
  return (A ** 2).sum()

A = jnp.array([[1.0, 2.0], [3.0, 4.0]])
df_dA = jax.grad(f)(A)
print(df_dA[0, 0], df_dA[0, 1], df_dA[1, 2])

个字符
或者，您可以将数组的条目拆分为单独的函数参数，然后使用argnums仅针对您感兴趣的那些进行梯度：

def f(a, b, c, d):
  A = jnp.array([[a, b], [c, d]])
  return (A ** 2).sum()

df_da, df_db, df_dc = jax.grad(f, argnums=(0, 1, 2))(1.0, 2.0, 3.0, 4.0)
print(df_da, df_db, df_dc)

2.0 4.0 8.0

的字符串
一般来说，你可能会发现第一种方法在实践中更容易使用，而且效率更高。它确实有一些浪费的计算，但坚持矢量化计算通常会带来净收益，特别是如果你在GPU或TPU等加速器上运行。