mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax
History
Jake VanderPlas 09d1b6d8d5 Deprecate jnp.msort following deprecation of numpy.msort
2022-12-07 10:08:18 -08:00
..
_src
Deprecate jnp.msort following deprecation of numpy.msort
2022-12-07 10:08:18 -08:00
example_libraries
Revert part of #13498
2022-12-05 09:21:09 -08:00
experimental
Fix type annotation for bcoo_update_layout.
2022-12-07 04:25:25 -08:00
image
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
interpreters
Remove typing_extensions dependency
2022-12-05 15:42:26 -08:00
lax
Add support for Hessenberg and tridiagonal matrix reductions on CPU.
2022-11-09 06:23:55 -08: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
Add support for Hessenberg and tridiagonal matrix reductions on CPU.
2022-11-09 06:23:55 -08:00
tools
Use the 11/09 libtpu build for jaxlib release since that passes all the tests.
2022-11-14 20:37:41 -08: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
Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.
2022-11-28 14:31:48 -08: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
Ensure that the initial dynamic_trace_state is canonicalized.
2022-12-06 20:39:53 -08: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
add a basic prototype of piles, behind jax_dynamic_shapes
2022-11-06 17:03:04 -08:00
monitoring.py
Add JAX monitoring library that instruments code via events.
2022-11-15 12:41:41 -08: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
Use jax.Array by default for doctests
2022-11-15 11:52:22 -08:00
sharding.py
Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh.
2022-12-01 19:28:32 -08: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
Finish jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release)
2022-11-15 09:02:53 -08:00