transformer_engine/jax/flax/transformer.py · be055eb027106afd3c5bd797c8617e509d11f0bb · 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

transformer.py 90.1 KB