Hotfix for compatibility with older jax versions - #1926

Merged

dpanici merged 19 commits into

dp/hotfix-batching

Sep 25, 2025

Commits on Sep 20, 2025

hotfix, add compatibility layer to sharding.update call as it is different api before jax 0.6.2
dpanici
committed
6fdb6e7
Revert "hotfix, add compatibility layer to sharding.update call as it is different api before jax 0.6.2"
dpanici
committed
1cac804
Revert "Work toward parallelization via sharding (multi-device) composed with vmap (across memory in single device) (#1869)"
dpanici
committed
de7e9e2

Commits on Sep 22, 2025

add compatibility layer
YigitElma
committed
3e98ed0
move the unrelated import outside
YigitElma
committed
579f9d2
add newer versions of JAX to CI workflow
YigitElma
committed
a67a0d7
remove broken jax versions
YigitElma
committed
5e7d27c

Commits on Sep 23, 2025

add specific version check, add newer jax versions to test
YigitElma
committed
9e1bb13
increase the cutoff version
YigitElma
committed
8a3b7e7
update dependency min/max versions
YigitElma
committed
5054bd0

Commits on Sep 24, 2025

add the jax version for the old version
YigitElma
committed
4aebb89
remove 0.7.2 for Kaya's sake
YigitElma
committed
46ae9cf
Merge branch 'master' into dp/hotfix-batching
YigitElma
authored
04dba80
Ku/batch (#1930)
unalmis
authored
41eec0e
update changelog
dpanici
committed
7cf2e7a
reduce test resolution
dpanici
committed
7d7cb19
Merge branch 'master' into dp/hotfix-batching
rahulgaur104
authored
172295a
Update tests/test_compute_funs.py
dpanici
authored
e8ea73f
Update tests/test_compute_funs.py

dpanici
and
github-actions[bot]
authored
b2f4c4a