mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-22 23:46:04 +00:00
Code Issues Packages Projects Releases Wiki Activity
d4bd2570ae
rocm_jax/tests
History
Adam Paszke d4bd2570ae [Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts 
PiperOrigin-RevId: 737956598
2025-03-18 04:47:51 -07:00
..
filecheck
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
mosaic
[Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts
2025-03-18 04:47:51 -07:00
notebooks
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
pallas
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
testdata
…
third_party/scipy
Fix CI test failure in line_search_test.
2024-05-15 07:38:20 -07:00
ann_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
aot_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
api_test.py
Raise an error if jax.config.update('jax_num_cpu_devices', val) is called after backend is initialized
2025-03-13 14:45:53 -07:00
api_util_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
array_api_skips.txt
array API: improve test coverage
2024-12-05 12:56:07 -08:00
array_api_test.py
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
array_interoperability_test.py
[dlpack] Support more DLPack dtypes now that we target DLPack 1.1
2025-03-14 09:10:56 -07:00
array_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
attrs_test.py
ref errors
2024-12-18 07:46:14 +00:00
batching_test.py
Change the way that batching.spec_types is updated.
2025-03-15 22:58:44 -07:00
blocked_sampler_test.py
Compute tile index using tile-based coordinates
2025-03-17 15:46:27 -07:00
BUILD
[JAX] Disable msan and asan for the profiler test running on nvidia gpu
2025-03-04 10:34:11 -08:00
cache_key_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
checkify_test.py
Don't monkey-patch functions in test_utils to count events for tests.
2024-12-12 09:58:14 -05:00
cholesky_update_test.py
Bump minimum jaxlib version to v0.4.30.
2024-06-18 12:35:08 -04:00
clear_backends_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
colocated_python_test.py
Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.
2025-03-06 14:57:52 -08:00
compilation_cache_test.py
#sdy unskip JAX Shardy tests that are already passing
2025-01-23 09:26:38 -08:00
config_test.py
Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.
2024-07-22 15:09:15 -07:00
core_test.py
shard_map and other fixes to direct-linearize
2025-03-07 21:02:40 +00:00
cudnn_fusion_test.py
Re-enable cudnn_fusion_test on A100.
2024-11-01 15:48:07 +00:00
custom_linear_solve_test.py
Handle effects in lax.custom_linear_solve.
2025-02-03 11:14:48 -05:00
custom_partitioning_sharding_rule_test.py
[jax:custom_partitioning] Support SdyShardingRule with multiple leading
2025-01-23 08:20:38 -08:00
custom_root_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
debug_info_test.py
A few more fixes for debug_info tests with direct_linearize.
2025-03-08 07:47:24 -05:00
debug_nans_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
debugger_test.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
debugging_primitives_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
deprecation_test.py
[internal] add deprecation test utilities
2025-01-10 11:54:09 -08:00
device_test.py
Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57
2024-07-29 18:44:31 -07:00
distributed_test.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
dtypes_test.py
Fix lint
2025-03-07 04:52:58 +00:00
dynamic_api_test.py
Port tests away from setUpClass and setUpModule to setUp alone.
2025-01-08 08:14:50 -08:00
error_check_test.py
Support error checking in explicit mode
2025-03-14 18:58:26 -07:00
errors_test.py
Merge pull request from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
experimental_rnn_test.py
Fix integer overflow in workspace size computations for experimental.rnn.*.
2025-03-12 08:22:04 -07:00
export_back_compat_test.py
#sdy Add JAX backwards compatibility test.
2025-03-03 06:01:34 -08:00
export_harnesses_multi_platform_test.py
[export] Fixes for export_harnesses_multi_platform_test.
2025-01-29 06:12:33 -08:00
export_test.py
Make the signature of AbstractMesh to be AbstractMesh(axis_size: tuple[int, ...], axis_name: tuple[str, ...], *, axis_types) instead of AbstractMesh(shape_tuple: tuple[tuple[str, int], ...], *, axis_types) so that we are consistent across all Mesh APIs: Mesh, AbstractMesh and make_mesh
2025-03-12 21:32:31 -07:00
extend_test.py
[better_errors] Make it explicit that debug_info is not None.
2025-02-13 22:07:04 -08:00
ffi_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
fft_test.py
Relax some test tolerances in N-D FFT tests.
2024-12-20 03:21:53 -08:00
for_loop_test.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
fused_attention_stablehlo_test.py
Removed unused import
2025-03-05 15:18:43 +00:00
garbage_collection_guard_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
generated_fun_test.py
Add a variant of safe_map() that has no return value, named foreach().
2025-03-14 07:42:48 -07:00
gpu_memory_flags_test.py
Add CUDA custom call example as a JAX test
2024-06-17 15:21:49 -04:00
heap_profiler_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
image_test.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
infeed_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
jax_jit_test.py
Add test case for jit compile of vmap on gpu.
2025-01-23 17:50:19 -08:00
jax_to_ir_test.py
[jax2tf] Remove non-native serialization test from jax_to_ir_test
2024-10-07 04:21:38 -07:00
jaxpr_effects_test.py
Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.
2025-02-21 09:45:59 -08:00
jaxpr_util_test.py
add pjit forwarding rule
2024-05-25 17:46:01 +00:00
jet_test.py
Adding support for copy_p primitive to jet.
2024-11-07 17:06:02 -08:00
key_reuse_test.py
[key reuse] fix signature for device_put
2025-01-17 09:47:50 -08:00
lax_autodiff_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
lax_control_flow_test.py
Increase the number of iterations in a test that compares rolled versus unrolled HLO for length.
2025-03-11 13:45:19 -07:00
lax_metal_test.py
adopt jax.extend.backend
2024-10-07 12:27:35 -07:00
lax_numpy_einsum_test.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
lax_numpy_indexing_test.py
Improved errors when indexing with floats
2025-02-28 15:04:07 -08:00
lax_numpy_operators_test.py
Implement jax.numpy.spacing
2024-10-03 10:40:39 -07:00
lax_numpy_reducers_test.py
jax.numpy reductions: avoid upcast of f16 when dtype is specified by user
2025-02-12 11:49:39 -08:00
lax_numpy_test.py
jax.numpy ndim/shape/size: deprecate non-array input
2025-03-04 10:42:32 -08:00
lax_numpy_ufuncs_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
lax_numpy_vectorize_test.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
lax_scipy_sparse_test.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
lax_scipy_special_functions_test.py
Update gammainc and gammaincc against scipy 1.16: return nan whenever one of operands is nan.
2025-03-12 17:48:45 +02:00
lax_scipy_spectral_dac_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_scipy_test.py
Merge pull request from mattjj:logsumexp-where-grad-nan
2025-01-10 11:27:35 -08:00
lax_test.py
add experimental lax.optimization_barrier autodiff rules
2025-03-14 22:40:55 +00:00
lax_vmap_op_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_vmap_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
layout_test.py
Make sure default layout is None for input and output layout in all codepaths
2025-02-27 14:26:25 -08:00
lazy_loader_test.py
…
linalg_sharding_test.py
Temporarily skip some more linalg sharding checks.
2025-02-28 12:26:18 -08:00
linalg_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
lobpcg_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
logging_test.py
logging fixes
2024-11-05 13:28:17 -08:00
lru_cache_test.py
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1
2025-02-14 12:01:55 -08:00
magma_linalg_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
memories_test.py
PR : Change the default value of print_operand_shape_ to false and print_large_constants_ to true.
2025-03-11 03:17:04 -07:00
mesh_utils_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
metadata_test.py
Disable test_source_file_prefix_removal test because there is cross-contamination of metadata information from different call sites because of cached jitted functions
2024-07-12 12:07:24 -07:00
mock_gpu_test.py
Annotate several tests as thread-unsafe.
2025-01-10 11:24:39 -08:00
mock_gpu_topology_test.py
Annotate several tests as thread-unsafe.
2025-01-10 11:24:39 -08:00
monitoring_test.py
Allow duration event listeners to take extra keyword arguments.
2023-10-04 13:15:49 -07:00
mosaic_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
multi_device_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
multibackend_test.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
multiprocess_gpu_test.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
mutable_array_test.py
Rename AxisTypes enum to AxisType
2025-03-14 11:48:21 -07:00
name_stack_test.py
Just use jit as the string in error messages instead of jit and pjit based on resource_env. This is to start deprecating the need for with mesh and replace it with use_mesh(mesh).
2025-03-05 20:09:30 -08:00
nn_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
ode_test.py
[XLA:CPU] Use central difference to calculate numerical gradient
2025-01-22 07:49:43 -08:00
optimizers_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
package_structure_test.py
C++ tree with path API
2024-11-30 21:26:48 -08:00
pgle_test.py
Fix test failure in PGLE tests.
2025-02-18 09:44:16 +00:00
pickle_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
pjit_test.py
Allow pspec to be passed to device_put if there is a mesh in the surrounding context
2025-03-17 17:47:56 -07:00
pmap_test.py
A few more fixes for debug_info tests with direct_linearize.
2025-03-08 07:47:24 -05:00
polynomial_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
pretty_printer_test.py
Compute source maps when pretty-printing jaxprs.
2024-05-06 15:45:25 -04:00
profiler_test.py
Add a profiler test for gpu run
2025-02-28 13:45:46 -08:00
python_callback_test.py
[xla:cpu] Implement XLA FFI handlers for CPU Jax callbacks.
2025-02-12 13:53:36 -08:00
pytorch_interoperability_test.py
Set PYTHONWARNINGS=error in bazel tests.
2024-09-24 12:30:11 -07:00
qdwh_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
ragged_collective_test.py
add autodiff rules for jax.lax.ragged_all_to_all collective
2025-03-11 18:22:02 -07:00
random_lax_test.py
Add replace option to random.categorical to enable sampling without replacement.
2025-03-17 13:41:46 -04:00
random_test.py
Don't monkey-patch functions in test_utils to count events for tests.
2024-12-12 09:58:14 -05:00
roofline_test.py
Support convolution in roofline.
2025-03-14 12:26:20 -07:00
scaled_matmul_stablehlo_test.py
Merge pull request from wenscarl:wenscarl/nvfp4
2025-03-13 13:30:46 -07:00
scipy_fft_test.py
Fix dtype bug in jax.scipy.fft.idct
2024-09-25 12:55:43 -04:00
scipy_interpolate_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_ndimage_test.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
scipy_optimize_test.py
Fix some warnings causing CI failures on ARM.
2024-09-24 17:25:26 -07:00
scipy_signal_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_spatial_test.py
test: change random matrix generation for Rotation
2025-01-31 09:21:46 -08:00
scipy_stats_test.py
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf
2025-02-01 12:51:11 +05:00
shape_poly_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
shard_alike_test.py
Use the mesh of out_aval when converting GSPMDSharding to NamedSharding. This makes sure that the axis types of the corresponding output is correct.
2025-02-20 17:13:24 -08:00
shard_map_test.py
Rename AxisTypes enum to AxisType
2025-03-14 11:48:21 -07:00
source_info_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
source_mapper_test.py
Merge pull request from justinjfu:sourcemap_win
2025-02-06 12:32:10 -08:00
sourcemap_test.py
Merge pull request from gspschmid:gschmid/sourcemaps
2024-05-08 05:44:08 -07:00
sparse_bcoo_bcsr_test.py
[JAX] disable flaky parameter permutations for sparse_bcoo_bcsr test.
2025-02-03 16:02:06 -08:00
sparse_nm_test.py
[XLA:GPU] Skip small tile sizes for sparse gemms on Ampere as well. Enable the JAX test again that has been failing.
2024-11-11 08:58:15 -08:00
sparse_test.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
sparsify_test.py
[better_errors] Improve error message for lax.switch branches output structure mismatch
2025-01-10 08:03:33 +02:00
stack_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
state_test.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
stax_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
string_array_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
svd_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
transfer_guard_test.py
Removed trivial docstrings from JAX tests
2025-03-17 07:49:37 -07:00
tree_util_test.py
Add an option to simplify keystr output and use a custom separator.
2025-01-21 10:18:42 -08:00
typing_test.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
util_test.py
Add a C++ implementation of a toplogical sort.
2025-03-14 16:04:25 -07:00
version_test.py
Modify version test to consider "rc" versions as well
2025-03-10 11:10:18 -07:00
warnings_util_test.py
Switch to a new thread-safe utility for catching warnings.
2025-01-09 11:58:34 -05:00
x64_context_test.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
xla_bridge_test.py
Use backend._get_all_devices() to validate devices.
2025-01-24 11:09:16 -08:00
xla_interpreter_test.py
Move some utilities out of dispatch.py next to their users, add more types.
2023-08-08 10:52:11 -07:00
xla_metadata_test.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