mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Adam Paszke a06a5aaa22 [Pallas] Use tpu.matmul instead of vector.contract in Mosaic lowering 
This will let us do mixed precision matmuls, which are rejected by the
vector.contract verifier.

PiperOrigin-RevId: 572901961
2023-10-12 07:39:19 -07:00
..
clusters
Fix typos across the package
2023-09-22 14:54:31 -07:00
debugger
jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186
2023-07-31 15:36:43 -07:00
extend
move wrap_key_data to jax.random
2023-09-18 14:38:22 -07:00
image
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
internal_test_util
Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().
2023-09-27 12:10:43 -07:00
interpreters
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
lax
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
lib
[CUDA] Add code to jax initialization that verifies that the CUDA libraries that are found are at least as new as the versions against which JAX was built.
2023-09-27 11:28:40 -07:00
nn
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
numpy
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
ops
MAINT Clean up leftover Array = Any aliases in jax/_src/**.py
2023-10-01 12:19:21 +01:00
pallas
[Pallas] Use tpu.matmul instead of vector.contract in Mosaic lowering
2023-10-12 07:39:19 -07:00
scipy
Merge pull request #17760 from superbobry:array-any
2023-10-03 08:50:07 -07:00
state
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
third_party
Add a type stub for jax.numpy.
2023-08-22 11:50:49 -07: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
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
ad_checkpoint.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
ad_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
api_util.py
Fix resolve_argnums if inspect.signature fails. In this case, donate_argnames was None leading an error in assert_no_intersection.
2023-07-31 12:10:15 -07:00
api.py
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
array.py
CI: update mypy to v1.6.0
2023-10-11 12:54:51 -07:00
basearray.py
Split basearray into separate Bazel module.
2023-03-13 11:14:41 -07:00
basearray.pyi
[typing] fix a few array type declarations
2023-09-12 13:21:48 -07:00
cache_key.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
callback.py
DOC: mention ShapeDtypeStruct in callback docstrings
2023-10-11 19:55:39 -07:00
checkify.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
cloud_tpu_init.py
Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().
2023-10-10 10:59:37 -07:00
compilation_cache_interface.py
Move implementation of compilation cache out of jax/experimental and into jax/_src.
2023-04-17 08:35:53 -07:00
compilation_cache.py
Set up plumbing for adding new compilation-cache-key generation algorithm.
2023-08-09 18:16:22 -07:00
compiler.py
Only perform compilation cache writes from process 0.
2023-10-09 13:55:07 -07:00
config.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
core.py
CI: update mypy to v1.6.0
2023-10-11 12:54:51 -07: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
Fix type failures under an upcoming pytype change.
2023-04-07 07:10:09 -07:00
custom_derivatives.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
custom_transpose.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
debugging.py
[JAX] Keep CPU host callbacks alive via IFRT, rather than by attaching them to the Python object.
2023-10-05 15:07:03 -07:00
deprecations.py
Make deprecation warnings warn at appropriate stacklevel
2023-05-23 14:43:38 -07:00
dispatch.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
distributed.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
dlpack.py
Remove take_ownership support in DLPack.
2023-10-10 09:43:02 -07:00
dtypes.py
Make jax._src.typing.DTypeLike more strictly defined
2023-10-10 09:01:19 -07:00
effects.py
[callbacks] Add support for shardable ordered effects.
2023-09-18 02:50:25 -07:00
environment_info.py
Split _src files custom_api_util, deprecations, effects and environment_info into separate Bazel targets.
2023-03-10 12:26:05 -08:00
errors.py
Fix typos across the package
2023-09-22 14:54:31 -07:00
flatten_util.py
Avoid imports from the public jax.* namespace in more places internally.
2023-04-04 11:41:40 -07:00
gfile_cache.py
Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads
2023-06-14 08:07:22 +08:00
jaxpr_util.py
Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target.
2023-08-08 10:09:09 -07:00
lax_reference.py
add an exp2 primitive and lax.exp2
2023-07-28 12:33:49 -07:00
lazy_loader.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
linear_util.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
logging_config.py
Add jax_debug_log_modules config option.
2023-07-28 18:11:12 +00:00
maps.py
Introduce a LoweringParameters dataclass for easier plumbing
2023-09-29 08:23:05 +03:00
mesh.py
Fix mypy error
2023-08-29 13:25:12 -07:00
monitoring.py
Allow duration event listeners to take extra keyword arguments.
2023-10-04 13:15:49 -07:00
op_shardings.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
partition_spec.py
Improve the sharding documentation.
2023-08-03 11:12:02 -04:00
path.py
Make the etils[epath] dependency of JAX optional.
2022-10-12 21:16:48 +00:00
pickle_util.py
Use cloudpickle for serializing host callbacks
2023-06-09 13:34:43 -07:00
pjit.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
pretty_printer.py
Update flags to use the ABSL typed flag API.
2023-07-27 12:15:58 -07:00
prng.py
[random] add shaped_abstractify handler for custom PRNG key
2023-10-10 16:15:19 -07:00
profiler.py
Fix typos across the package
2023-09-22 14:54:31 -07:00
public_test_util.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
random.py
Make jax._src.typing.DTypeLike more strictly defined
2023-10-10 09:01:19 -07:00
sharding_impls.py
Make is_fully_addressable an abstract method and implement it on each concrete Sharding.
2023-09-05 17:28:22 -07:00
sharding_specs.py
Replace references to DeviceArray with Array.
2023-08-18 17:46:00 -04:00
sharding.py
Make is_fully_addressable an abstract method and implement it on each concrete Sharding.
2023-09-05 17:28:22 -07:00
source_info_util.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
stages.py
Migrate another subset of internal modules to use state objects
2023-10-11 08:46:06 -07:00
test_util.py
test: prevent regressions on object comparisons
2023-10-06 14:23:14 -07:00
tpu_custom_call.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
traceback_util.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
tree_util.py
fast dispatch for functions over typed PRNG key arrays
2023-09-13 09:43:58 -07:00
typing.py
Make jax._src.typing.DTypeLike more strictly defined
2023-10-10 09:01:19 -07:00
util.py
Migrate a subset of internal modules to use state objects
2023-10-09 07:29:53 -07:00
xla_bridge.py
Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().
2023-10-10 10:59:37 -07:00