mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 12:46:08 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Sharad Vikram a6c4c87f3e Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
..
clusters
Don't set TPU topology env vars in cloud_tpu_init.py
2023-01-13 22:02:16 +00:00
debugger
Replace uses of jnp.ndarray with jax.Array inside JAX.
2023-02-15 14:53:00 -08:00
image
migrate more internal dependencies from jax.core to jax._src.core
2023-02-14 23:01:11 -08:00
internal_test_util
lax_test: Create a separate module for lax-specific test utils in a new package.
2023-02-16 15:29:41 -08:00
interpreters
Merge pull request #14549 from sharadmv:dbidx-effects
2023-02-17 23:43:38 -08:00
lax
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
lib
MeshComputation.cost_analysis() isn't implemented with PJRT C API.
2023-02-16 00:15:24 +00:00
nn
migrate more internal dependencies from jax.core to jax._src.core
2023-02-14 23:01:11 -08:00
numpy
jnp.ndarray.view: implement all dtypes
2023-02-17 10:54:37 -08:00
ops
migrate more internal dependencies from jax.core to jax._src.core
2023-02-14 23:01:11 -08:00
scipy
DOC: fix typo in sph_harm
2023-02-16 15:34:18 +01:00
state
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
third_party
Address drastic slowdown in mypy runtime
2022-12-16 14:48:26 -08:00
__init__.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
abstract_arrays.py
Expose fp8 in jax dtypes and mlir builder.
2023-01-13 18:12:12 -08:00
ad_checkpoint.py
Refactor effects system to use effect types, not objects
2023-02-17 17:40:08 -08:00
ad_util.py
Merge pull request #13693 from jakevdp:typing-ad-util
2022-12-20 16:50:11 -08:00
api_util.py
Make error_test a jax_test so that we can test other configs and fix it with jit/pjit merge.
2023-01-17 18:43:05 -08:00
api.py
Merge pull request #14549 from sharadmv:dbidx-effects
2023-02-17 23:43:38 -08:00
array.py
[Rollforward] Convert _arrays to return PyArray instead of PyBuffer.
2023-02-17 11:52:43 -08:00
basearray.py
TMP
2023-02-17 09:49:11 -08:00
basearray.pyi
[typing] define additional methods & properties on jax.Array
2023-02-10 09:42:32 -08:00
callback.py
Refactor effects system to use effect types, not objects
2023-02-17 17:40:08 -08:00
checkify.py
Refactor effects system to use effect types, not objects
2023-02-17 17:40:08 -08:00
cloud_tpu_init.py
Add environment variable check in additional to libtpu check for cloud tpu vm
2023-01-25 09:51:37 -08:00
config.py
Remove uses of jax.xla_computation from metadata_test.py
2023-02-15 17:26:21 -08:00
core.py
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
custom_api_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
custom_batching.py
migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching
2023-02-09 15:11:57 -08:00
custom_derivatives.py
Refactor effects system to use effect types, not objects
2023-02-17 17:40:08 -08:00
custom_transpose.py
migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad
2023-02-06 22:52:36 -08:00
debugging.py
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
deprecations.py
Deprecate jax.interpreters.xla.Device and jax.interpreters.xla.DeviceArray.
2023-02-09 16:11:48 -08:00
device_array.py
DOC: improve documentation for jax.Array methods
2023-02-14 13:04:27 -08:00
dispatch.py
Merge pull request #14549 from sharadmv:dbidx-effects
2023-02-17 23:43:38 -08:00
distributed.py
Add Open MPI automatic distributed initialization
2023-01-11 17:08:09 -06:00
dlpack.py
[Rollforward] Convert _arrays to return PyArray instead of PyBuffer.
2023-02-17 11:52:43 -08:00
dtypes.py
DOC: add docs for jax.dtypes module
2023-02-14 11:18:59 -08:00
effects.py
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
environment_info.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
errors.py
migrate internal dependencies from jax.core to jax._src.core
2022-12-16 21:00:14 -08:00
flatten_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
global_device_array.py
Return jax.Array from GDA's callback APIs if jax.Array is True.
2023-02-16 16:02:05 -08:00
iree.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
lax_reference.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
lazy_loader.py
Avoid loading scipy eagerly.
2022-10-12 19:51:09 +00:00
linear_util.py
Move jax.linear_util to jax._src.linear_util
2022-12-20 14:49:27 -08:00
maps.py
Merge pull request #14549 from sharadmv:dbidx-effects
2023-02-17 23:43:38 -08:00
monitoring.py
Move contents of jax.monitoring to jax._src.monitoring.
2023-02-08 15:03:22 -08:00
path.py
Make the etils[epath] dependency of JAX optional.
2022-10-12 21:16:48 +00:00
pjit.py
Add JaxprInputEffect and refactor StateEffects to use it
2023-02-21 16:30:06 -08:00
pretty_printer.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
prng.py
Rename jax.sharding.OpShardingSharding to jax.sharding.GSPMDSharding. jax.sharding.OpShardingSharding will be removed in 3 months from Feb 17, 2023.
2023-02-17 17:11:06 -08:00
profiler.py
Add a number of missing function cross-references in the docs.
2022-11-07 12:00:26 -05:00
public_test_util.py
Roll-back https://github.com/google/jax/pull/14144 due to downstream test failures
2023-01-25 12:15:36 -08:00
random.py
migrate internal dependencies from jax.interpreters.batching to jax._src.interpreters.batching
2023-02-09 15:11:57 -08:00
sharding.py
Replace op_sharding_sharding with gspmd_sharding. This is purely an internal change.
2023-02-17 17:53:13 -08:00
source_info_util.py
Let is_user_frame ignore frames from stdlib.
2022-11-08 10:50:08 -08:00
stages.py
Remove more exported names from jax.interpreters.xla.
2023-02-16 11:56:30 -08:00
test_util.py
Add clear_cache endpoint to python pjit and cpp pjit functions.
2023-02-14 18:46:25 -08:00
traceback_util.py
add blank line, mainly to trigger/test source sync
2023-02-01 13:56:29 -08:00
tree_util.py
Use jnp.arange to break the pjit cache (when jit and pjit are merged) because pytest runs tests non-hermetically.
2023-02-08 10:17:37 -08:00
typing.py
Remove typing_extensions dependency
2022-12-05 15:42:26 -08:00
util.py
[Rollforward] Convert _arrays to return PyArray instead of PyBuffer.
2023-02-17 11:52:43 -08:00