mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-21 14:26:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
Mathew Odden 6b35155294
Fix invalid lowerings for ROCm in Pallas (#223) 
popcount and clz were effectively broken on ROCm,
since math_dialect had incorrect lowerings.

Use the device intrinsics for these functions, as
well as for exp and absf, which fixes some accuracy issues in
the pallas tests.

Docs for OCML/OCKL

- https://github.com/ROCm/llvm-project/blob/amd-staging/amd/device-libs/doc/OCML.md
- https://github.com/ROCm/llvm-project/blob/amd-staging/amd/device-libs/doc/OCKL.md
2025-02-14 11:27:52 -06:00
..
clusters
Adds a flag to control proxy env checking.
2024-11-16 00:00:59 -08:00
cudnn
support head size of 256
2024-12-19 18:38:06 +00:00
debugger
Cleanup: fix unused imports & mark exported names
2024-10-16 17:42:41 -07:00
export
Activate FFI implementation of tridiagonal reduction on GPU.
2025-01-10 09:28:15 -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
Activate FFI implementation of tridiagonal reduction on GPU.
2025-01-10 09:28:15 -08:00
interpreters
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -08:00
lax
[sharding_in_types] Error out for reshape for splits like this: (4, 6, 8) -> (4, 4, 2, 6)
2025-01-17 06:58:29 -08:00
lib
[mosaic_gpu] Added a serialization pass
2025-01-17 03:12:51 -08:00
nn
[sharding_in_types] Functions like einsum, reshape, broadcast_in_dim, broadcasted_iota, convert_element_type and sharding_cast that take out_sharding as an argument in their signature should also allow PartitionSpec instead of just NamedSharding as an input.
2025-01-08 11:11:16 -08:00
numpy
[sharding_in_types] Rename .at[...].get(out_spec) to .at[...].get(out_sharding).
2025-01-16 18:56:52 -08:00
ops
fix grad(logsumexp) to produce 0s where where is False
2025-01-08 23:38:06 +00:00
pallas
Fix invalid lowerings for ROCm in Pallas (#223)
2025-02-14 11:27:52 -06:00
scipy
enable partitionable threefry by default
2025-01-13 22:46:24 -08:00
state
[better_errors] More cleanup
2025-01-15 10:22:29 +00:00
third_party
feat(lib): add real-valued implementation of jax.scipy.special.fresnel
2024-09-03 09:50:19 +02:00
__init__.py
…
abstract_arrays.py
Internal: use a single registry for abstractify APIs
2024-12-23 08:44:35 -08:00
ad_checkpoint.py
[better_errors] More cleanup
2025-01-15 10:22:29 +00: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] More cleanup
2025-01-15 10:22:29 +00:00
api.py
Merge pull request #25827 from gnecula:debug_info_2
2025-01-14 09:12:37 -08:00
array.py
Allow resharding between tokens on a single device
2025-01-16 11: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
#sdy enable pure callbacks and debug prints in JAX.
2025-01-09 13:37:51 -08:00
checkify.py
[better_errors] More cleanup
2025-01-15 10:22:29 +00: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
[AutoPGLE] Explicitly ignore host callback pointers
2024-11-26 04:06:15 -08:00
compiler.py
[AutoPGLE] FIx PGLE kokoro test failures.
2025-01-07 08:59:59 -08:00
compute_on.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00
config.py
enable partitionable threefry by default
2025-01-13 22:46:24 -08:00
core.py
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -08:00
custom_api_util.py
…
custom_batching.py
[better_errors] More cleanup
2025-01-15 10:22:29 +00:00
custom_derivatives.py
add mutable array ref error checks to cond and custom_vjp
2024-12-20 01:44:50 +00:00
custom_partitioning_sharding_rule.py
Implement the extension to the custom_partitioning API.
2025-01-08 13:34:47 -08:00
custom_partitioning.py
[better_errors] More cleanup
2025-01-15 10:22:29 +00:00
custom_transpose.py
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.
2024-10-29 11:04:31 -07:00
debugging.py
#sdy enable pure callbacks and debug prints in JAX.
2025-01-09 13:37:51 -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
Make result_handler of _DeferredShardArg a method instead of a property. Also play some code golf.
2025-01-16 09:53:48 -08:00
distributed.py
Add more input validation to jax.distributed.initialize.
2024-12-19 16:26:34 -08:00
dlpack.py
Deprecated calling `jax.dlpack.from_dlpack` with a DLPack tensor
2024-09-03 15:16:02 -07:00
dtypes.py
Support int4/uint4 in jnp.ndarray.view
2024-12-20 13:57:40 -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
Add a register_custom_type_id function to the GPU plugins.
2025-01-07 07:29:38 -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
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04: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] More cleanup
2025-01-15 10:22:29 +00: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
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -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
[sharding_in_types] Enforce AxisTypes to always exist if set_mesh is used.
2024-12-10 18:03:21 -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
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -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 float8_e4m3 and float8_e3m4 types support
2024-11-08 18:58:31 +00:00
random.py
fix gamma_p in vmap-based impl rule mode
2024-12-13 15:32:09 -08:00
shard_alike.py
Remove internal uses of api_util.shaped_abstractify
2024-12-19 07:06:36 -08:00
sharding_impls.py
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -08:00
sharding_specs.py
Fix type annotations for NumPy 2.2
2024-12-11 14:24:58 -08:00
sharding.py
[jax] Typing on common_devices_indices_map
2024-12-02 12:40:36 -08:00
source_info_util.py
Use a class with __slots__ instead of a NamedTuple in JaxprEqn and SourceInfo, which are two tuples we build frequently.
2024-07-01 19:18:58 -07:00
sourcemap.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
stages.py
Jax: Stop returning a list of cost-analyses.
2025-01-15 09:53:59 -08:00
test_util.py
Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
2025-01-16 17:55:54 -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
[Mosaic TPU] Append dump id to timestamp to make dump list ordered
2025-01-14 12:44:10 -08:00
traceback_util.py
Run pyupgrade --py310-plus.
2024-06-26 16:10:18 -04:00
tree_util.py
Bump xla_extension_version after jaxlib release.
2024-12-18 08:07:19 -05: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
Make utils for reporting function name work with functools.partial by using the inner .func attribute if the object doesn't have a __name__ attribute. functools.partial objects do not have __name__ attributes by default.
2025-01-15 11:40:59 -08:00
xla_bridge.py
[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.
2025-01-08 06:37:44 -08:00
xla_metadata.py
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
2024-11-05 08:32:25 -08:00