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/tests
History
wenscarl 638c6ae046 Add e8m0fnu support by conditional dtype.
2025-01-22 21:57:43 +00:00
..
filecheck
Register acos primitive to lower to CHLO acos.
2024-10-25 13:20:36 -07:00
mosaic
[Mosaic GPU] Add layout inference for splat arith.ConstantOps and vector.SplatOps.
2025-01-16 07:18:35 -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 GPU] Add manual consumed barrier handling to WS pipeline.
2025-01-22 10:59:58 -08:00
testdata
…
third_party/scipy
…
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] Refactor debug info tests
2025-01-20 20:21:01 +01:00
api_util_test.py
…
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
Make make_mesh take visible_axes, hidden_axes and collective_axes as parameters instead of axis_types to make it a more cleaner API.
2025-01-21 09:02:39 -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
…
BUILD
[better_errors] Expand the tests for debug_info
2025-01-22 16:49:16 +01: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
…
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
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
config_test.py
…
core_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
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] Make SdyShardingRule a user facing class.
2024-12-10 13:05:43 -08:00
custom_root_test.py
…
debug_info_test.py
[better_errors] Expand the tests for debug_info
2025-01-22 16:49:16 +01: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
Annotate several tests as thread-unsafe.
2025-01-10 11:24:39 -08:00
deprecation_test.py
[internal] add deprecation test utilities
2025-01-10 11:54:09 -08:00
device_test.py
…
dtypes_test.py
Add e8m0fnu support by conditional dtype.
2025-01-22 21:57:43 +00: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
Merge pull request #25755 from ROCm:ci_rnn_final-upstream
2025-01-15 10:40:54 -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
Port tests away from setUpClass and setUpModule to setUp alone.
2025-01-08 08:14:50 -08:00
export_test.py
Port tests away from setUpClass and setUpModule to setUp alone.
2025-01-08 08:14:50 -08:00
extend_test.py
Move jex.ffi to jax.ffi.
2024-12-29 13:06:19 +00:00
ffi_test.py
[internal] add deprecation test utilities
2025-01-10 11:54:09 -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
…
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
Remove internal uses of api_util.shaped_abstractify
2024-12-19 07:06:36 -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
…
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
Add support for constants in the decomposition of lax.composites.
2025-01-21 13:28:42 -08:00
lax_vmap_op_test.py
…
lax_vmap_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
layout_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
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
…
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
…
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
[internal] add deprecation test utilities
2025-01-10 11:54:09 -08: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
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
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
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
pytorch_interoperability_test.py
Set PYTHONWARNINGS=error in bazel tests.
2024-09-24 12:30:11 -07:00
qdwh_test.py
…
random_lax_test.py
Don't monkey-patch functions in test_utils to count events for tests.
2024-12-12 09:58:14 -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
Don't perform size 0 slices into scipy rotations.
2024-11-08 09:56:12 -05: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
[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
shard_map_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
source_info_test.py
…
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
…
transfer_guard_test.py
…
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
Switch to a new thread-safe utility for catching warnings.
2025-01-09 11:58:34 -05: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