Revision 2cf102c2de8fec5b492b411d6c915acb65a1821a authored by Peter Hawkins on 06 April 2023, 21:19:53 UTC, committed by jax authors on 07 April 2023, 00:12:27 UTC
Move dependencies of sharding_impls into Bazel targets lower in the dependency graph: * Move ParsedPartitionSpec and ArrayMapping into partition_spec.py. This breaks a dependencies on pjit.py and pxla.py * Move AxisEnv out of xla.py into mesh.py. * Move AxisContext and friends out of mlir.py into mesh.py. This removes the need for sharding_impls.py to depend on mlir.py, flattening the dependency graph. * Move prepare_axis_resources into sharding_impls.py out of pjit.py. * Fix a handful of new pytype errors. PiperOrigin-RevId: 522435255
1 parent c7b99e6
| File | Mode | Size |
|---|---|---|
| .github | ||
| benchmarks | ||
| build | ||
| cloud_tpu_colabs | ||
| docs | ||
| examples | ||
| images | ||
| jax | ||
| jaxlib | ||
| tests | ||
| third_party | ||
| .bazelrc | -rw-r--r-- | 14.7 KB |
| .bazelversion | -rw-r--r-- | 6 bytes |
| .gitignore | -rw-r--r-- | 342 bytes |
| .pre-commit-config.yaml | -rw-r--r-- | 993 bytes |
| .readthedocs.yml | -rw-r--r-- | 569 bytes |
| AUTHORS | -rw-r--r-- | 313 bytes |
| CHANGELOG.md | -rw-r--r-- | 84.4 KB |
| CITATION.bib | -rw-r--r-- | 408 bytes |
| CONTRIBUTING.md | -rw-r--r-- | 150 bytes |
| LICENSE | -rw-r--r-- | 11.1 KB |
| README.md | -rw-r--r-- | 25.7 KB |
| WORKSPACE | -rw-r--r-- | 1.6 KB |
| conftest.py | -rw-r--r-- | 2.5 KB |
| mypy.ini | -rw-r--r-- | 1.3 KB |
| pylintrc | -rw-r--r-- | 1.6 KB |
| pyproject.toml | -rw-r--r-- | 90 bytes |
| pytest.ini | -rw-r--r-- | 2.0 KB |
| setup.cfg | -rw-r--r-- | 1.5 KB |
| setup.py | -rw-r--r-- | 5.4 KB |
Computing file changes ...
