1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import pytest
import awkward as ak
jax = pytest.importorskip("jax")
def test():
ak.jax.register_and_check()
array = ak.Array([[1.0, 2.0, 3.0], [], [4.0, 5.0]], backend="jax")
val_mean, grad_mean = jax.value_and_grad(ak.mean, argnums=0)(array)
_, grad_sum = jax.value_and_grad(ak.sum, argnums=0)(array)
assert val_mean == 3
assert ak.all(
grad_mean == ak.Array([[0.2, 0.2, 0.2], [], [0.2, 0.2]], backend="jax")
)
|
