mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
Adam Paszke 85c6b6a128 [Mosaic GPU] Add support for tiling stores to refs using small tiling 
The difficulty here is that our register tiling is based on the (64, 8)
shape, while the memory tiling is now (8, swizzle // bytewidth). Before,
we would assume that each register tile fits neatly within a single
memory tile, but now it is obviously not the case. Luckily, it wasn't
too hard to add.

PiperOrigin-RevId: 734517000
2025-03-07 05:19:11 -08:00
..
filecheck
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
mosaic
[Mosaic GPU] Add support for tiling stores to refs using small tiling
2025-03-07 05:19:11 -08: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
[Mosaic TPU] Support fp8 upcast to f32
2025-03-06 17:19:15 -08: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
[aot] Add support for as_text(debug_info=True).
2025-01-10 07:59:56 +02:00
api_test.py
Dedup non-ref constants closed in cond branch functions.
2025-03-07 04:01:42 -08:00
api_util_test.py
Consolidate handling of input argument resolution in custom_* APIs.
2024-08-13 00:30:23 -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
Increase the minimum NumPy version to v1.25.
2024-12-18 08:18:57 -05:00
array_test.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
attrs_test.py
ref errors
2024-12-18 07:46:14 +00:00
batching_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
blocked_sampler_test.py
[Pallas] Implement block-invariant sampling.
2024-06-24 11:20:39 -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
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -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
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08: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
[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.
2025-02-25 14:43:17 +01:00
debug_nans_test.py
Fix shard_map debug_nan leakage of manual out_avals in the impl rules of jit i.e. impl rule of jit saw a manual out_aval which is not expected. This is a band-aid for now with a TODO to do a proper fix
2025-02-24 10:15:21 -08: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
Internal: avoid adding _DimExpr to dtypes._weak_types
2025-01-29 09:11:02 -08: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
Minor bug fixes in error checking
2025-03-06 06:57:52 -08:00
errors_test.py
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
experimental_rnn_test.py
Update RNN kernels to use FFI.
2025-02-06 18:27:58 -08: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
#sdy support JAX export tests when Shardy is enabled.
2025-03-03 04:57:06 -08: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
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
generated_fun_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -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
Update references to JAX's GitHub repo
2025-03-04 18:31:09 -08: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
Make logaddexp and logaddexp2 into ufuncs
2024-11-18 09:27:36 -08: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 jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.
2024-11-18 20:33:27 -05: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 #25791 from mattjj:logsumexp-where-grad-nan
2025-01-10 11:27:35 -08:00
lax_test.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
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
Fix linalg.norm to return zero for proper norms of empty matrices.
2025-03-03 15:02:34 -05: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
Remove the skip for test_output_streaming_inside_scan
2025-03-03 14:54:03 -08:00
mesh_utils_test.py
[Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh
2024-11-28 09:24:20 -08: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
…
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
[better_errors] Cleanup use of DebugInfo.arg_names and result_paths
2025-02-23 08:27:56 +02: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
Improve after review # 2
2025-02-26 04:48:25 +00:00
ode_test.py
[XLA:CPU] Use central difference to calculate numerical gradient
2025-01-22 07:49:43 -08:00
optimizers_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
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
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
pjit_test.py
Auto complete specs in a sharding if aval.ndim > len(sharding.spec) with None. So that for a 2D input, P('data') continues to work.
2025-03-06 16:10:14 -08:00
pmap_test.py
[better_errors] Refactor debug info tests
2025-01-20 20:21:01 +01: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
Bump up tolerance for tests using _testQdwh.
2024-06-14 23:18:20 -07:00
ragged_collective_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
random_lax_test.py
Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such
2025-03-02 19:42:40 -08: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
roofline: Support computing flops for unary ops.
2025-03-06 17:44:36 -08:00
scaled_matmul_stablehlo_test.py
Add apache header.
2025-02-26 15:26:56 -06: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
Make api_test.py work when test cases are run using multiple threads.
2025-01-08 14:09:07 -08: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
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
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 #26369 from justinjfu:sourcemap_win
2025-02-06 12:32:10 -08:00
sourcemap_test.py
Merge pull request #21096 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
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -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
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
string_array_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
svd_test.py
Bump some SVD test tolerances.
2024-07-15 12:54:00 -07:00
transfer_guard_test.py
fix some additional warnings related to #21834
2024-06-13 16:06:14 -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
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
version_test.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
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
…
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