mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax
History
Yash Katariya cc5af7ed98 Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding. 
`MeshPspecSharding` name will be replaced with `NamedSharding` in 3 months.

PiperOrigin-RevId: 485753078
2022-11-02 19:13:13 -07:00
..
_src
Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.
2022-11-02 19:13:13 -07:00
example_libraries
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
experimental
Merge pull request #13083 from jakevdp:jet-dynamic-update
2022-11-02 17:44:09 -07:00
image
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
interpreters
Add a new jax.spmd_mode config for preventing unintentional hangs and incorrect results when users pass jax.Arrays that span across multiple processes (i.e. not fully addressable) to jit or jnp operations (that are jitted by default).
2022-10-31 09:51:42 -07:00
lax
Add lax.cumlogsumexp for cumulative logsumexp operations.
2022-10-31 15:08:52 -07:00
lib
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
nn
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
numpy
Point to ndarray.at from docstring of unimplemented jnp.put & jnp.place
2022-10-28 14:13:36 -07:00
ops
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
scipy
implement truncnorm in jax.scipy.stats
2022-10-22 15:48:20 -04:00
tools
Update Colab TPU driver version
2022-10-11 17:49:10 -07:00
__init__.py
Add a new jax.spmd_mode config for preventing unintentional hangs and incorrect results when users pass jax.Arrays that span across multiple processes (i.e. not fully addressable) to jit or jnp operations (that are jitted by default).
2022-10-31 09:51:42 -07:00
abstract_arrays.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ad_checkpoint.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
api_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
BUILD
Move array.py and sharding.py from experimental/ to _src/.
2022-09-27 10:06:52 -07:00
cloud_tpu_init.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
collect_profile.py
Cleanup: remove unused imports in private modules
2022-10-20 14:37:21 -07:00
config.py
Cleanup: remove unused imports in private modules
2022-10-20 14:37:21 -07:00
core.py
Merge pull request #13022 from mattjj:leak-checker-improvements
2022-10-28 16:05:43 -07:00
custom_batching.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
custom_derivatives.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
custom_transpose.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
debug.py
Add inspect_array_sharding, enabling looking at shardings in pjit-ted functions
2022-09-22 17:36:56 -07:00
distributed.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
dlpack.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
dtypes.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
errors.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
flatten_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
jaxpr_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
linear_util.py
Remove error-prone most_recent_entry() support from lu.cache.
2022-10-27 16:41:44 -07:00
prng.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
profiler.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
py.typed
Improve type annotations
2021-01-13 10:26:35 -05:00
random.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
sharding.py
Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding.
2022-11-02 19:13:13 -07:00
stages.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
test_util.py
Remove deprecated functionality from jax.test_util
2022-10-11 08:16:34 -07:00
tree_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
version.py
Increase the minimum jaxlib version to 0.3.22.
2022-10-27 10:24:11 -04:00