mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 13:26:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax
History
Jake VanderPlas 05faf0f40d Remove deprecated functionality from jax.test_util 
PiperOrigin-RevId: 480360504
2022-10-11 08:16:34 -07:00
..
_src
Merge pull request #12727 from jakevdp:scipy-linalg-types
2022-10-11 06:34: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
[jax2tf] Implement jax2tf(pjit) for experimental_native_lowering
2022-10-11 09:45:07 +02:00
image
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
interpreters
Support MANUAL collectives in top-level xmaps
2022-10-11 06:47:30 -07:00
lax
Rolling forward https://github.com/google/jax/pull/12707 after rollback, due to changes in relatively trivial jax.numpy shape validation code failed in some downstream user tests.
2022-10-10 18:51:37 -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
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
ops
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
scipy
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
tools
Pin default jax.tools.colab_tpu.setup_tpu driver version.
2022-09-30 17:45:49 +00:00
__init__.py
Move array.py and sharding.py from experimental/ to _src/.
2022-09-27 10:06:52 -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
Fix collect_profile _src import
2022-10-04 14:31:17 -07:00
config.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
core.py
Rolling forward https://github.com/google/jax/pull/12707 after rollback, due to changes in relatively trivial jax.numpy shape validation code failed in some downstream user tests.
2022-10-10 18:51:37 -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
Rolling forward https://github.com/google/jax/pull/12707 after rollback, due to changes in relatively trivial jax.numpy shape validation code failed in some downstream user tests.
2022-10-10 18:51:37 -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
Move array.py and sharding.py from experimental/ to _src/.
2022-09-27 10:06:52 -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
Update version and CHANGELOG for jax 0.3.21 release
2022-10-03 22:06:19 +00:00