tests/jax/test_helper.py · a7a1a07024cca5535edab97ca93531de4a93149f · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Move jax.experimental.maps.Mesh to jax.sharding.Mesh (#276) · df6f347f
zlsh80826 authored Jun 14, 2023
```
Move jax.experimental.maps.Mesh to jax.sharding.Mesh
Signed-off-by: Reese Wang <rewang@nvidia.com>
```
df6f347f

test_helper.py 9.8 KB