1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import pytest
import awkward as ak
numba = pytest.importorskip("numba")
jax = pytest.importorskip("jax")
def test():
behavior = {}
ak.jax.register_and_check()
input_arr = ak.Array([1.0], backend="jax")
@numba.vectorize(
[
numba.float32(numba.float32, numba.float32),
numba.float64(numba.float64, numba.float64),
]
)
def _some_kernel(x, y):
return x * x + y * y
@ak.mixin_class(behavior)
class SomeClass:
@property
def some_kernel(self):
return _some_kernel(self.x, self.y)
ak.behavior.update(behavior)
arr = ak.zip({"x": input_arr, "y": input_arr}, with_name="SomeClass")
assert ak.all(arr.some_kernel == ak.Array([2.0], backend="jax"))
|
