mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
tttc3 b1b56ea0b0 Enable pivoted QR on GPU via MAGMA. 
Originally noted in #20282, this commit provides a GPU compatible
implementation of `geqp3` via MAGMA.
2025-02-12 16:12:42 +00:00
..
clusters
Adds a flag to control proxy env checking.
2024-11-16 00:00:59 -08:00
cudnn
disable head dim 256 on bw now
2025-02-05 22:07:48 +00:00
debugger
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
export
Remove old backward compatibility mode for old PRGN custom call on GPU
2025-02-04 07:34:52 -08:00
extend
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
image
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
internal_test_util
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
interpreters
[better_errors] Continue adding debug info to Jaxprs (step 8)
2025-02-12 14:23:52 +01:00
lax
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
lib
[pallas:triton] The lowering now uses PTX instead of Triton IR
2025-02-03 13:21:40 -08:00
nn
Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg
2025-01-31 15:41:50 -08:00
numpy
Merge pull request #26447 from jakevdp:refactor-contractions
2025-02-12 07:14:30 -08:00
ops
refactor: move scalar type defs out of lax_numpy.py
2025-02-06 14:48:10 -08:00
pallas
[Pallas/Mosaic GPU] Enable progressive lowering for integer addition.
2025-02-12 06:29:25 -08:00
scipy
Enable pivoted QR on GPU via MAGMA.
2025-02-12 16:12:42 +00:00
state
[better_errors] Continue adding debug info to Jaxprs (step 6)
2025-02-11 11:28:58 +01: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] Add a canonicalize_value step before dispatching bind so that we can insert mesh_casts under the following conditions:
2025-02-03 18:00:19 -08:00
ad_checkpoint.py
[better_errors] Add debug info to the Jaxprs formed for AD
2025-02-05 19:21:02 +02:00
ad_util.py
Cleanup: remove uses of no-op raise_to_shaped
2024-12-12 09:49:06 -08:00
api_util.py
[better_errors] Continue adding debug info to Jaxprs (step 6)
2025-02-11 11:28:58 +01:00
api.py
[better_errors] Continue adding debug info to Jaxprs (step 4)
2025-02-08 09:13:55 +02:00
array.py
Return arrays from ArrayImpl._check_and_rearrange.
2025-02-05 09:24: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
Add "sequential_unrolled" vmap method for callbacks.
2025-02-11 06:09:16 -08:00
checkify.py
[better_errors] Continue adding debug info to Jaxprs (step 6)
2025-02-11 11:28:58 +01:00
cloud_tpu_init.py
Fix KeyError recently introduced in cloud_tpu_init.py
2024-11-20 17:46:06 +00:00
compilation_cache_interface.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
compilation_cache.py
Fix cache init when JAX Array is created early (#25768)
2025-01-22 03:44:29 +00:00
compiler.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
compute_on.py
Optimize implementation of the compute_on context manager.
2025-01-26 09:24:33 -08:00
config.py
Make JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES env vars
2025-01-28 17:17:56 -08:00
core.py
Merge pull request #26384 from gnecula:debug_info_jaxpr_4
2025-02-10 07:42:57 -08:00
custom_api_util.py
…
custom_batching.py
[better_errors] Continue adding debug info to Jaxprs (step 5)
2025-02-07 08:23:10 +02:00
custom_dce.py
[better_errors] Add debug info to more Jaxprs and WrappedFun (step 1)
2025-02-04 10:02:35 +02:00
custom_derivatives.py
[better_errors] Continue adding debug info to Jaxprs (step 6)
2025-02-11 11:28:58 +01: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
[better_errors] Continue adding debug info to Jaxprs (step 8)
2025-02-12 14:23:52 +01:00
custom_transpose.py
[better_errors] Continue adding debug info to Jaxprs (step 4)
2025-02-08 09:13:55 +02:00
debugging.py
Revert https://github.com/jax-ml/jax/pull/25982 since callbacks can now use JAX functions.
2025-01-29 11:12:32 -05: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
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
distributed.py
improve runtime error message
2025-01-29 21:03:26 +00:00
dlpack.py
Pulled DLDeviceType to XLA backend mapping into a global
2025-01-29 11:38:50 -08:00
dtypes.py
refactor: move array creation routines out of lax_numpy.py
2025-02-06 15:47:30 -08:00
earray.py
Internal: use a single registry for abstractify APIs
2024-12-23 08:44:35 -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
errors.py
Consolidate material on PRNGs and add a short summary to Key Concepts.
2024-11-15 14:44:57 -08:00
ffi.py
Match output container to result_shape_dtypes in ffi_call.
2025-02-11 17:33:32 -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
Enable the enhanced TPU launch barrier on all TPU generations.
2024-09-11 15:28:14 -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] Continue adding debug info to Jaxprs (step 8)
2025-02-12 14:23:52 +01:00
logging_config.py
Adding JAX_LOGGING_LEVEL configuration option
2024-11-05 09:56:46 -08:00
lru_cache.py
Bumped jaxlib version mypy uses on the CI
2024-07-26 11:22:39 +01:00
mesh_utils.py
Fix type annotations for NumPy 2.2
2024-12-11 14:24:58 -08:00
mesh.py
Speed up NamedSharding construction.
2025-02-07 12:20:51 -08:00
monitoring.py
Add JAX events that have time spans, not only durations.
2025-01-07 23:08:30 -08:00
op_shardings.py
Fix type annotations for NumPy 2.2
2024-12-11 14:24:58 -08:00
partition_spec.py
Add a custom __reduce__ for UnconstrainedSingleton because it can be picked and then loaded back and we need the id of P.UNCONSTRAINED to match before and after loading.
2025-02-11 20:05:48 -08:00
path.py
Import etils.epath lazily
2024-08-13 14:48:38 -07:00
pickle_util.py
Removed noop # type: ignore comments
2024-05-19 21:01:29 +01:00
pjit.py
[sharding_in_types] Error out if the sharding's specs passed to with_sharding_constraint don't refer to Auto axes.
2025-02-11 10:16:52 -08: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
Fix pmap with sharded typed prng key
2025-01-10 18:20:09 -08:00
profiler.py
Small fix for the jax trace dumping path
2024-08-26 10:51:08 -07:00
public_test_util.py
add docstrings for check_vjp and check_jvp
2025-01-23 17:31:13 -08:00
random.py
Roll back multinomial change from https://github.com/jax-ml/jax/pull/25688
2025-02-05 09:13:56 -08:00
shard_alike.py
Remove internal uses of api_util.shaped_abstractify
2024-12-19 07:06:36 -08:00
sharding_impls.py
Use an PartitionSpec.UNCONSTRAINED to represent unconstrained dimensions in ParsedPartitionSpec, rather than None.
2025-02-09 08:45:01 -08:00
sharding_specs.py
Fix type annotations for NumPy 2.2
2024-12-11 14:24:58 -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
Speed up name stack printing.
2025-01-22 16:24:16 -08:00
sourcemap.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
stages.py
docstrings and API reference doc listing for the traced AOT stage
2025-02-11 22:30:50 -08:00
test_util.py
Remove code for compatibility with jaxlib < 0.5.0.
2025-02-06 13:08:32 -08: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
Add an option to simplify keystr output and use a custom separator.
2025-01-21 10:18:42 -08: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
Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.
2025-02-06 14:30:36 -08:00
xla_metadata.py
Fix CI failures in xla_metadata_test.
2025-01-30 17:28:54 -08:00