tests/jax/test_distributed_layernorm_mlp.py · 6673f1658c605f7999e04a43bde1493c4b9ad741 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Flax with compute dtype inferred from input dtype. (#1485) · 6673f165
Phuong Nguyen authored Feb 18, 2025
```
flax module with compute dtype inferred from the inputs
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
```
6673f165

test_distributed_layernorm_mlp.py 12.7 KB