=================================== FAILURES ===================================
___________________________ test_jax_transform_warn ____________________________
[gw40] linux -- Python 3.13.11 /nix/store/slhpx9glq7vl99bwi93bgrhn3syv98s1-python3-3.13.11/bin/python3.13
getkey = GetKey(seed=9044076, call=11, key=Array([ 0, 9044076], dtype=uint32))
def test_jax_transform_warn(getkey):
class A(eqx.Module):
linear: Callable
class B(eqx.Module):
linear: Callable
def __init__(self, linear):
self.linear = linear
for cls in (A, B):
for transform in (
jax.jit,
jax.grad,
jax.vmap,
jax.value_and_grad,
jax.jacfwd,
jax.jacrev,
jax.hessian,
jax.custom_jvp,
jax.custom_vjp,
jax.checkpoint, # pyright: ignore
jax.pmap,
):
> with pytest.warns(
match="Possibly assigning a JAX-transformed callable as an attribute"
):
E Failed: DID NOT WARN. No warnings of type (<class 'Warning'>,) were emitted.
E Emitted warnings: [].
tests/test_module.py:885: Failed
=========================== short test summary info ============================
FAILED tests/test_module.py::test_jax_transform_warn - Failed: DID NOT WARN. No warnings of type (<class 'Warning'>,) were emitted.
================== 1 failed, 457 passed, 1 skipped in 24.91s ===================
Here is a failure we noticed when using
jax>=0.9.0.Context: NixOS/nixpkgs#482118