Thanks!
@jit
def block0(a, b, n):
# transonic block (
# float a, b;
# int n
# )
result = 0.0
for _ in range(n):
There is also a third error which would require supporting fori_loop
from the lower level lax
library. This can be marked as an expected failure for now.
https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#fori-loop
This is lovely! Thanks a lot! After a first quick read, the code seems reasonable. I'm going to investigate about the 2 failures.
Ashwin V. Mohanan (463565d4) at 14 Mar 21:34
Generate data_tests with jax backend
... and 3 more commits
Ashwin V. Mohanan (214bc651) at 14 Mar 21:33
Generate data_tests with jax backend
... and 3 more commits
FAILED tests/test_meson.py::test_meson_option - AssertionError: assert False
FAILED tests/test_run.py::test_create_pythran_files - AssertionError: assert False
These tests are failing, and I am not sure how to fix them.
Fixes #57
Ashwin V. Mohanan (46d54135) at 14 Mar 21:24
Generate data_tests with jax backend
... and 3 more commits
Oh yes it would be very nice to add a JAX backend!
I guess it should be quite simple.
First one needs to add a transonic/backends/jax.py
file with something very similar to what is in
https://foss.heptapod.net/fluiddyn/transonic/-/blob/branch/default/src/transonic/backends/numba.py (less than 100 lines of code).
Tiny change here https://foss.heptapod.net/fluiddyn/transonic/-/blob/branch/default/src/transonic/backends/__init__.py#L39
Fix + add tests!
In the simplest terms, to support Jax we have to
@jit
decorator like in Numba andimport numpy as np
with import jax.numpy as np
The API has of course some deviations, but for many cases this should work. Shall we add it here, alongside the Numba backend implementation?
Pierre Augier (4e01b62d) at 20 Feb 12:58
Merge branch 'topic/default/init-meson' into 'branch/default'
Pierre Augier (60650ae8) at 20 Feb 12:58
Added tag 0.6.3 for changeset 75da913d79da
Pierre Augier (bbd47512) at 20 Feb 12:57
Pierre Augier (4e01b62d) at 20 Feb 12:57
Merge branch 'topic/default/init-meson' into 'branch/default'
... and 5 more commits
Pierre Augier (4685e8cc) at 20 Feb 08:30
Doc transonic-init-meson
Pierre Augier (608034b8) at 20 Feb 05:13
Prepare 0.6.3
Pierre Augier (ecf4e11e) at 20 Feb 05:09
cmd transonic-init-meson