examples/jax/encoder/test_model_parallel_encoder.py · 5b6ef0545ed14497ff296d18c3618ace4ba391db · 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_model_parallel_encoder.py 16.2 KB