mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expose get_abstract_mesh via the jax.sharding namespace

Browse Source 
PiperOrigin-RevId: 736972976
This commit is contained in:
Yash Katariya 2025-03-14 13:44:32 -07:00 committed by jax authors
parent a11d8891ce
commit aa9480a441
3 changed files with 700 additions and 701 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/notebooks/explicit-sharding.ipynb
View File

@ -62,9 +62,8 @@
"import jax\n",
"import numpy as np\n",
"import jax.numpy as jnp\n",
"from jax.sharding import PartitionSpec as P, AxisType, set_mesh\n",
"from jax.sharding import PartitionSpec as P, AxisType, set_mesh, get_abstract_mesh\n",
"from jax.experimental.shard import reshard, auto_axes\n",
"from jax._src.mesh import get_abstract_mesh\n",
"\n",
"jax.config.update('jax_num_cpu_devices', 8)"
]

3
docs/notebooks/explicit-sharding.md
View File

@ -59,9 +59,8 @@ outputId: a64bcbcb-27f8-4c57-8931-8091c9bb8ebf
import jax
import numpy as np
import jax.numpy as jnp
from jax.sharding import PartitionSpec as P, AxisType, set_mesh
from jax.sharding import PartitionSpec as P, AxisType, set_mesh, get_abstract_mesh
from jax.experimental.shard import reshard, auto_axes
from jax._src.mesh import get_abstract_mesh
jax.config.update('jax_num_cpu_devices', 8)
```

3
jax/sharding.py
View File

@ -28,10 +28,11 @@ from jax._src.sharding_impls import (
from jax._src.partition_spec import (
PartitionSpec as PartitionSpec,
)
from jax._src.interpreters.pxla import Mesh as Mesh
from jax._src.mesh import (
Mesh as Mesh,
AbstractMesh as AbstractMesh,
AxisType as AxisType,
get_abstract_mesh as get_abstract_mesh,
)
_deprecations = {