jax

7bdcd36c - Add a test which shows that fsdp matmul with unreduced works without custom_vjp