mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax
History
Yash Katariya bcd4048dd5 Set the mesh of tangent.aval when we are creating zeros_like_aval because when you close over an array which is unused, we error out during canonicalization 
PiperOrigin-RevId: 729340808
2025-02-20 19:32:34 -08:00
..
_src
Set the mesh of tangent.aval when we are creating zeros_like_aval because when you close over an array which is unused, we error out during canonicalization
2025-02-20 19:32:34 -08:00
example_libraries
update example optimizers library docstring
2024-09-03 23:40:47 -07:00
experimental
[pallas:mosaic_gpu] Added support for binary/comparison ops with WG semantics
2025-02-20 15:06:27 -08:00
extend
[better_errors] Make it explicit that debug_info is not None.
2025-02-13 22:07:04 -08:00
image
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
interpreters
Merge pull request #26078 from gnecula:debug_info_jaxpr
2025-02-04 10:54:26 -08:00
lax
Reorder top-level functions in lax.linalg, and add/expand docstrings.
2025-02-13 12:57:55 -08:00
lib
Finalize deprecation of some symbols from jax.lib.xla_client
2024-12-23 10:14:16 -08:00
nn
Remove a number of expired deprecations.
2024-10-31 15:40:54 -07:00
numpy
Merge pull request #26401 from jakevdp:numpy-consts
2025-02-18 11:32:25 -08:00
ops
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
scipy
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf
2025-02-01 12:51:11 +05:00
tools
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
2025-02-05 10:01:23 -08:00
__init__.py
Expose get_ty aka get_aval from jax namespace
2025-02-18 21:22:19 -08:00
ad_checkpoint.py
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
api_util.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
BUILD
[pallas:mosaic_gpu] Added support for binary/comparison ops with WG semantics
2025-02-20 15:06:27 -08:00
cloud_tpu_init.py
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
collect_profile.py
Use pathlib for profiler log_dir
2024-07-28 18:47:45 +03:00
core.py
[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo
2025-02-02 06:23:03 +02:00
custom_batching.py
Use PEP484-style exports in several submodules
2024-08-14 08:59:56 -07:00
custom_derivatives.py
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
custom_transpose.py
Use PEP484-style exports in several submodules
2024-08-14 08:59:56 -07:00
debug.py
…
distributed.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
dlpack.py
Use PEP484-style exports in several submodules
2024-08-14 08:59:56 -07:00
dtypes.py
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
errors.py
Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.
2024-09-26 08:39:30 -07:00
export.py
[export] Add support for serialization for some custom PyTree nodes
2024-10-21 11:38:13 +02:00
ffi.py
[xla:python] Add a mechanism for "batch partitioning" of FFI calls.
2025-02-07 09:14:06 -08:00
flatten_util.py
…
monitoring.py
Add JAX events that have time spans, not only durations.
2025-01-07 23:08:30 -08:00
profiler.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
py.typed
…
random.py
Roll back multinomial change from https://github.com/jax-ml/jax/pull/25688
2025-02-05 09:13:56 -08:00
sharding.py
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -08:00
stages.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
test_util.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
tree_util.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
tree.py
Add jax.tree shortcuts for .*_with_path calls, for convenience of users.
2024-12-12 15:13:32 -08:00
typing.py
Fix typo in jax.typing module doc
2024-09-11 23:34:03 +10:00
util.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
version.py
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
2025-02-05 10:01:23 -08:00