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

Move jax.interpreters.pxla to jax._src.interpreters.pxla.

Browse Source 
Make jax.interpreters.pxla a shim that at the moment re-exports everything in the implementation, with the goal of reducing it over time.

PiperOrigin-RevId: 507584264
This commit is contained in:
Peter Hawkins 2023-02-06 14:28:36 -08:00 committed by jax authors
parent 3d9ae6b467
commit 38a59a313b
5 changed files with 4203 additions and 3935 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/api.py
View File

@ -81,7 +81,7 @@ from jax.custom_transpose import custom_transpose
from jax.interpreters import partial_eval as pe
from jax.interpreters import mlir
from jax.interpreters import xla
from jax.interpreters import pxla
from jax._src.interpreters import pxla
from jax.interpreters import ad
from jax.interpreters import batching

3947
jax/_src/interpreters/pxla.py Normal file
View File

File diff suppressed because it is too large Load Diff

2
jax/_src/stages.py
View File

@ -594,7 +594,7 @@ class Lowered(Stage):
def compile(self) -> Compiled:
"""Compile, returning a corresponding ``Compiled`` instance."""
from jax.interpreters import pxla
from jax._src.interpreters import pxla
if (jax.config.jax_array and
isinstance(self._lowering, pxla.MeshComputation)):

4186
jax/interpreters/pxla.py
View File

File diff suppressed because it is too large Load Diff

1
setup.cfg
View File

@ -36,6 +36,7 @@ per-file-ignores =
jax/errors.py:F401
jax/flatten_util.py:F401
jax/interpreters/ad.py:F401
jax/interpreters/pxla.py:F401
jax/linear_util.py:F401
jax/prng.py:F401
jax/profiler.py:F401