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/_src
History
Yash Katariya 76dec38286 Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env. 
This would also make it easier to deprecate the `with mesh: pjit` path in the future from user code since the new path would be completely tested.
This will also allow us to remove `resource_env` from JAX and the internal API access of `resource_env.physical_mesh` spread throughout codebases internally and externally.

PiperOrigin-RevId: 735602187
2025-03-10 20:21:02 -07:00
..
clusters
Adds a flag to control proxy env checking.
2024-11-16 00:00:59 -08:00
cudnn
Merge pull request #26627 from Cjkkkk:remove_fmha_rewriter
2025-03-05 05:20:25 -08:00
debugger
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
export
support e2m1fn
2025-03-05 17:44:34 +00:00
extend
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
image
Migrated to mypy 1.14.1 with --allow_redefinition
2025-02-13 15:38:28 +00:00
internal_test_util
#sdy Add JAX backwards compatibility test.
2025-03-03 06:01:34 -08:00
interpreters
Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.
2025-03-10 20:21:02 -07:00
lax
Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.
2025-03-10 12:25:22 -07:00
lib
[Mosaic GPU] Remove old jaxlib version guards.
2025-02-12 08:49:40 -08:00
nn
Merge pull request #26765 from Qwlouse:patch-1
2025-03-04 08:30:45 -08:00
numpy
Merge pull request #26525 from wenscarl:e2m1fn
2025-03-10 11:46:18 -07:00
ops
doc: fix description of logsumexp axis
2025-02-28 12:53:33 -08:00
pallas
[Pallas TPU] Remove next_slot SMEM tensor from pipeline emitter
2025-03-10 17:19:39 -07:00
scipy
Prevent tracer leaks in scipy.special.expn
2025-03-06 14:38:11 -08:00
state
fixups that we meant to include in #26427
2025-03-08 00:03:26 +00:00
third_party
refactor: move scalar type defs out of lax_numpy.py
2025-02-06 14:48:10 -08:00
__init__.py
…
abstract_arrays.py
[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset. Also some more changes as listed below:
2025-02-12 10:03:01 -08:00
ad_checkpoint.py
Rename dot_with_no_batch_dims_saveable to dots_with_no_batch_dims_saveable for internal consistency
2025-03-10 13:04:49 -07:00
ad_util.py
Cleanup: remove uses of no-op raise_to_shaped
2024-12-12 09:49:06 -08:00
api_util.py
Internal: make it easier to detect the vmap sentinel
2025-03-10 11:37:50 -07:00
api.py
shard_map and other fixes to direct-linearize
2025-03-07 21:02:40 +00:00
array.py
Use batched_device_put for token shard_arg handler
2025-02-27 11:30:22 -08:00
basearray.py
Make Tracers non-hashable
2024-11-05 09:08:33 -08:00
basearray.pyi
[sharding_in_types] Rename .at[...].get(out_spec) to .at[...].get(out_sharding).
2025-01-16 18:56:52 -08:00
blocked_sampler.py
Remove internal KeyArray alias
2024-11-20 10:30:12 -08:00
cache_key.py
Remove dead code after minimum jaxlib version bump to v0.4.36.
2024-12-09 07:35:05 -08:00
callback.py
Reverts 6e83de5909b5dad6827c1682726c840c9688b32d
2025-02-21 18:34:16 -08:00
checkify.py
Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.
2025-03-10 20:21:02 -07:00
cloud_tpu_init.py
Warn the user if transparent huge pages aren't enabled.
2025-03-10 10:37:58 -07:00
compilation_cache_interface.py
Fix detection of epath
2025-03-04 11:44:27 -05:00
compilation_cache.py
Add JAX_COMPILATION_CACHE_EXPECT_PGLE option
2025-02-06 08:19:45 +00:00
compiler.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
compute_on.py
Minor format cleanup
2025-02-15 04:56:27 -08:00
config.py
Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such
2025-03-02 19:42:40 -08:00
core.py
shard_map and other fixes to direct-linearize
2025-03-07 21:02:40 +00:00
custom_api_util.py
…
custom_batching.py
Improve error messages when input argument resolution fails in custom_* APIs.
2025-03-04 10:31:35 -05:00
custom_dce.py
Improve error messages when input argument resolution fails in custom_* APIs.
2025-03-04 10:31:35 -05:00
custom_derivatives.py
Reverts 342cb7b99a09180472823a33c7cdad8a8db77875
2025-03-05 10:22:40 -08:00
custom_partitioning_sharding_rule.py
[jax:custom_partitioning] Support SdyShardingRule with multiple leading
2025-01-23 08:20:38 -08:00
custom_partitioning.py
Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.
2025-03-10 20:21:02 -07:00
custom_transpose.py
[better_errors] Continue adding debug info to Jaxprs (step 4)
2025-02-08 09:13:55 +02:00
debugging.py
Remove _parsed_pspec from everywhere in JAX except for NamedSharding constructor. I'll do that in the next CL since that has a dependency on C++ so needs guards.
2025-02-27 10:17:06 -08:00
deprecations.py
Add sph_harm_y to jax.scipy.special and deprecate sph_harm.
2025-01-09 12:53:00 -05:00
dispatch.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
distributed.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
dlpack.py
Pulled DLDeviceType to XLA backend mapping into a global
2025-01-29 11:38:50 -08:00
dtypes.py
Fix lint
2025-03-07 04:52:58 +00:00
earray.py
PRNGKeyArray.aval should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.
2025-03-01 18:18:19 -08:00
effects.py
Effects is now an immutable set
2023-12-11 08:45:52 -08:00
environment_info.py
Update device info printed in print_environment_info to print the TPU kind also along with global and local device count.
2024-10-04 09:20:27 -07:00
error_check.py
Better error message when raise_if_error() is called within a traced context
2025-03-10 16:55:06 -07:00
errors.py
Consolidate material on PRNGs and add a short summary to Key Concepts.
2024-11-15 14:44:57 -08:00
ffi.py
Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.
2025-02-21 09:45:59 -08:00
flatten_util.py
Avoid calls to warnings.catch_warnings in JAX core code.
2024-12-20 15:43:03 -05:00
hardware_utils.py
Warn the user if transparent huge pages aren't enabled.
2025-03-10 10:37:58 -07:00
jaxpr_util.py
[better_errors] Expand the tests for debug_info
2025-01-22 16:49:16 +01:00
lax_reference.py
Add int4/uint4 support to bitcast_convert_type
2024-12-20 12:45:24 -08:00
layout.py
sharding cleanup: use inline checks for unimplemented and auto
2024-10-25 04:22:40 -07:00
lazy_loader.py
lazy_loader.attach now only imports the submodule once on first access
2024-08-28 15:37:28 -07:00
linear_util.py
[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.
2025-02-25 14:43:17 +01:00
logging_config.py
Adding JAX_LOGGING_LEVEL configuration option
2024-11-05 09:56:46 -08:00
lru_cache.py
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1
2025-02-14 12:01:55 -08:00
mesh_utils.py
Reverts 0f0d5e90ef1c3d60f35020141710ea350d17816b
2025-02-27 13:27:32 -08:00
mesh.py
shard_map and other fixes to direct-linearize
2025-03-07 21:02:40 +00:00
monitoring.py
Add JAX events that have time spans, not only durations.
2025-01-07 23:08:30 -08:00
named_sharding.py
Merge _check_mesh_resource_axis and _check_axis_type_consistency into 1 function.
2025-02-27 12:51:25 -08:00
op_shardings.py
Merge pull request #26518 from superbobry:maint-2
2025-02-13 15:44:19 -08:00
partition_spec.py
Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec. We need to do this after sharding-in-types to speed up NamedSharding construction and remove a lot of tech debt and unnecessary complexity.
2025-02-27 08:59:25 -08:00
path.py
Fix detection of epath
2025-03-04 11:44:27 -05:00
pickle_util.py
Removed noop # type: ignore comments
2024-05-19 21:01:29 +01:00
pjit.py
Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.
2025-03-10 20:21:02 -07:00
pretty_printer.py
Make api_test.py work when test cases are run using multiple threads.
2025-01-08 14:09:07 -08:00
prng.py
PRNGKeyArray.aval should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.
2025-03-01 18:18:19 -08:00
profiler.py
Add JAX_COMPILATION_CACHE_EXPECT_PGLE option
2025-02-06 08:19:45 +00:00
public_test_util.py
support e2m1fn
2025-03-05 17:44:34 +00:00
random.py
Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such
2025-03-02 19:42:40 -08:00
shard_alike.py
Remove internal uses of api_util.shaped_abstractify
2024-12-19 07:06:36 -08:00
sharding_impls.py
#sdy support JAX export tests when Shardy is enabled.
2025-03-03 04:57:06 -08:00
sharding_specs.py
Merge pull request #26518 from superbobry:maint-2
2025-02-13 15:44:19 -08:00
sharding.py
[sharding_in_types] Add vmap + explicit sharding support. The main changes are:
2025-01-29 09:34:27 -08:00
source_info_util.py
[direct-linearize] fix name_stack bugs
2025-03-08 01:51:19 +00:00
sourcemap.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
stages.py
[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset. Also some more changes as listed below:
2025-02-12 10:03:01 -08:00
test_util.py
support e2m1fn
2025-03-05 17:44:34 +00:00
test_warning_util.py
Switch to a new thread-safe utility for catching warnings.
2025-01-09 11:58:34 -05:00
tpu_custom_call.py
[JAX/Pallas] Add has_side_effect parameter to CompilerParams to stop CSE of operations.
2025-02-11 13:33:01 -08:00
traceback_util.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
tree_util.py
Removed unused `# type: ignore` comments
2025-02-13 21:12:27 +00: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
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
util.py
Remove device_context from trace_context because we don't need it there. We can get compilation cache misses (and tracing/lowering cache hit) naturally without putting concrete devices into trace_context.
2025-01-21 16:21:36 -08:00
xla_bridge.py
Check "jax_rocm_visible_devices" at client creation.
2025-02-25 09:59:59 +01:00
xla_metadata.py
Fix CI failures in xla_metadata_test.
2025-01-30 17:28:54 -08:00