examples/jax/encoder/test_model_parallel_encoder.py · 5b16807c7b0fb9cd3ee3cf9f716e64271a071551 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Use keyword args for jit in_shardings and out_shardings (#1898) · 5b16807c
jberchtold-nvidia authored Jun 26, 2025
```
Use keyword args for jit in_shardings and out_shardings
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
```
5b16807c

test_model_parallel_encoder.py 19.8 KB