mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
carlosgmartin 32411a430f Add jax.random.multinomial.
2025-01-31 18:45:55 -05:00
..
filecheck
Register acos primitive to lower to CHLO acos.
2024-10-25 13:20:36 -07:00
mosaic
[mgpu] Correct instruction for conversion of unsigned int types.
2025-01-31 09:06:40 -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
Re-enable Pallas bf16 exp2 tests on TPU
2025-01-31 08:36:57 -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
…
aot_test.py
[aot] Add support for as_text(debug_info=True).
2025-01-10 07:59:56 +02:00
api_test.py
[better_errors] Add more debug info test coverage
2025-01-26 08:12:29 +02: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
Reverts bb951136e9b91a584bb422119ada76cc69c86024
2025-01-31 14:42:22 -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
Fix lax.ragged_all_to_all degenerate case
2025-01-30 12:05:02 -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
…
colocated_python_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
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
Disable two tests of GC behavior if there are multiple threads per process.
2025-01-31 09:14:49 -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
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
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
…
debug_info_test.py
Disable a debug_info_test test that fails in CI.
2025-01-30 17:25:18 -08:00
debug_nans_test.py
Add a temporary fix for spurious debug_nans errors when round-tripping jaxprs.
2024-10-30 11:34:08 -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
Always dispatch CPU executables synchronously when they include callbacks.
2025-01-28 18:23:35 -08: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
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
errors_test.py
Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.
2024-09-26 08:39:30 -07:00
experimental_rnn_test.py
Reverts 86643a1b3e0516e1a2ddbdabbb714cf8c0301f18
2025-01-31 08:05:46 -08:00
export_back_compat_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05: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
[export] Fix mis-used of NamedSharding in export tests
2025-01-24 09:18:02 +02:00
extend_test.py
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
ffi_test.py
Lower np.ndarray to DenseElementsAttr instead of ArrayAttr.
2025-01-31 11:06:06 -08: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
support head size of 256
2024-12-19 18:38:06 +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
…
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
…
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
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -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
[better_errors] Refactor debug info tests
2025-01-20 20:21:01 +01: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
Indexing: avoid dynamic_slice when mode='clip'
2025-01-14 11:20:50 -08:00
lax_numpy_operators_test.py
Implement jax.numpy.spacing
2024-10-03 10:40:39 -07:00
lax_numpy_reducers_test.py
[internal] add deprecation test utilities
2025-01-10 11:54:09 -08:00
lax_numpy_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05: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
…
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
Lower np.ndarray to DenseElementsAttr instead of ArrayAttr.
2025-01-31 11:06:06 -08:00
lax_vmap_op_test.py
…
lax_vmap_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
layout_test.py
[XLA:CPU] Add CPU client support for layout modes.
2025-01-30 07:27:02 -08:00
lazy_loader_test.py
…
linalg_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -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
Storing the last access time of a cache entry in a separate file
2024-07-03 13:54:36 +04:00
magma_linalg_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
memories_test.py
Fix remat bug on primitives with multiple outputs.
2025-01-14 10:26:58 -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
…
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
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
mutable_array_test.py
add mutable array ref error checks to cond and custom_vjp
2024-12-20 01:44:50 +00:00
name_stack_test.py
[aot] Add support for as_text(debug_info=True).
2025-01-10 07:59:56 +02:00
nn_test.py
fix error tolerance
2025-01-22 22:00:16 +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
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
pickle_test.py
…
pjit_test.py
Do automatic casting from auto -> manual when the context mesh is manual and avals are in auto mode. This happens when values are being closed over in a shard_map. The casting is happening at lax level but we can move this to a different place later on.
2025-01-30 13:14:04 -08:00
pmap_test.py
[better_errors] Refactor debug info tests
2025-01-20 20:21:01 +01:00
polynomial_test.py
…
pretty_printer_test.py
…
profiler_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
python_callback_test.py
Unconditionally skip async deadlock test for pure_callback.
2025-01-29 09:49:01 -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
Fix lax.ragged_all_to_all degenerate case
2025-01-30 12:05:02 -08:00
random_lax_test.py
Add jax.random.multinomial.
2025-01-31 18:45:55 -05: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
[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
scipy_fft_test.py
Fix dtype bug in jax.scipy.fft.idct
2024-09-25 12:55:43 -04:00
scipy_interpolate_test.py
…
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
…
scipy_spatial_test.py
test: change random matrix generation for Rotation
2025-01-31 09:21:46 -08:00
scipy_stats_test.py
Fix false positive debug_nans error caused by NaNs that are properly handled in jax.scipy.stats.gamma
2024-11-21 10:33:29 -08: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
[sharding_in_types] Upstream changes from defaulting sharding_in_types config to True experiment. There aren't a lot of failures in TGP but we can atleast upstream these changes until we work on the failures.
2025-01-28 11:04:42 -08:00
shard_map_test.py
#sdy Enable test_partial_auto_of_random_keys under Shardy.
2025-01-28 15:36:52 -08:00
source_info_test.py
…
source_mapper_test.py
Disable source mapper test for optimized hlo
2025-01-30 14:07:54 -08:00
sourcemap_test.py
…
sparse_bcoo_bcsr_test.py
Set PYTHONWARNINGS=error in bazel tests.
2024-09-24 12:30:11 -07: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
…
state_test.py
[better_errors] Improvements in propagation of debugging info
2025-01-20 15:09:51 +01: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
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
…
util_test.py
Stop using generators for linear_util transformations.
2024-11-13 13:47:07 -08:00
version_test.py
Adds an env that can let users provide a custom version suffix for jax dev build.
2024-11-16 14:39:54 -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