jax.experimental.sparse.value_and_grad#

jax.experimental.sparse.value_and_grad(fun, argnums=0, has_aux=False, **kwargs)[源代码]#

支持稀疏感知的 jax.value_and_grad() 版本

参数和返回值与 jax.value_and_grad() 相同，但当针对 jax.experimental.sparse 数组计算梯度时，梯度将在由数组稀疏模式定义的子空间中计算。

示例

>>> from jax.experimental import sparse
>>> X = sparse.BCOO.fromdense(jnp.arange(6.))
>>> y = jnp.ones(6)
>>> sparse.value_and_grad(lambda X, y: X @ y)(X, y)
(Array(15., dtype=float32), BCOO(float32[6], nse=5))
参数:

  • fun (可调用对象)

  • argnums (int | Sequence[int])

返回类型:

Callable[…, tuple[Any, Any]]