mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
jax authors 0e62c4cfcc Merge pull request #20976 from gnecula:export_fix_symtab 
PiperOrigin-RevId: 629125592
2024-04-29 11:14:04 -07:00
..
filecheck
tests: access tree utilities via jax.tree.*
2024-02-26 14:17:18 -08:00
mosaic
[Mosaic GPU] Add the flash attention example
2024-04-29 09:31:30 -07:00
notebooks
Deprecate arr.device_buffer and arr.device_buffers
2023-12-06 10:20:29 -08:00
pallas
Reverts 7844bac5d220b41253495cacf719f61905f46925
2024-04-29 11:13:43 -07:00
testdata
[PJRT C API] Add parsing PJRT client create options from json file.
2023-03-21 16:57:34 -07:00
third_party/scipy
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
ann_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
aot_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
api_test.py
Fix donation with kwargs. The problem is that pytrees sort dictionaries by default. So if we create the donation vector with original kwargs order, it won't match the aval order (which is created by sorting kwargs i.e. dict) and we end up donating the wrong input.
2024-04-23 14:50:04 -07:00
api_util_test.py
Replace donation_vector's logic with donation_vector_with_in_tree which is now deleted
2024-04-23 17:38:30 -07:00
array_api_test.py
Expose tile function in array_api namespace
2024-04-26 17:34:58 +00:00
array_interoperability_test.py
[PJRT C API] Add stream extension to support DLPack and implement this extension in CUDA plugin.
2024-04-19 10:41:55 -07:00
array_test.py
If callback returns a fully replicated global array, return it as is.
2024-04-14 14:35:57 -07:00
attrs_test.py
[attrs] fix tracer lifetime bug, fixes #20082
2024-03-05 12:08:44 -08:00
batching_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
BUILD
Disable sparse_nm_test_gpu_h100 because of flakiness
2024-04-26 13:13:15 -07:00
cache_key_test.py
Strip device_assignment on GPU platform.
2024-03-19 01:50:39 -07:00
checkify_test.py
MAINT Migrate remaining internal/test modules to use state objects
2023-10-12 17:32:15 +01:00
clear_backends_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
compilation_cache_test.py
Adds meaningful function names for better debugging.
2024-04-23 09:45:57 -07:00
core_test.py
[key reuse] rename flag to jax_debug_key_reuse
2024-03-22 05:37:30 -07:00
custom_linear_solve_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
custom_object_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
custom_root_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
debug_nans_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
debugger_test.py
jax.debug.callback now passes arguments as jax.Arrays
2024-04-19 13:57:18 -07:00
debugging_primitives_test.py
jax.debug.callback now passes arguments as jax.Arrays
2024-04-19 13:57:18 -07:00
deprecation_test.py
Fix tests that mistakenly used assertRaises(..., msg=...) to match a message.
2023-10-06 08:07:49 -04:00
dtypes_test.py
Added int4 and uint4 to dtype-specific tests
2024-04-18 15:20:20 +01:00
dynamic_api_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
errors_test.py
MAINT Migrate remaining internal/test modules to use state objects
2023-10-12 17:32:15 +01:00
experimental_rnn_test.py
MAINT Do not import the config object in JAX internals
2023-10-18 10:55:13 +01:00
export_back_compat_test.py
Reverts 9db5e693ebb4ad786c6e52b562cf32aeaba2e7e1
2024-04-26 04:14:34 -07:00
export_harnesses_multi_platform_test.py
[ROCm]: Add rocm as platform in export_harnessess_multi_platform_test
2024-03-20 16:24:06 +00:00
export_test.py
Fix mlir.merge_mlir_modules to properly remember the inlined symbols
2024-04-29 19:56:41 +03:00
extend_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
fft_test.py
Updated
2024-02-05 18:01:48 -05:00
for_loop_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
fused_attention_stablehlo_test.py
Support BNTH input formats
2024-04-03 20:48:37 +00:00
generated_fun_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
gfile_cache_test.py
Move implementation of compilation cache out of jax/experimental and into jax/_src.
2023-04-17 08:35:53 -07:00
gpu_memory_flags_test.py
Exclude test_gpu_memory_allocation from pytest execution.
2023-12-20 14:53:33 -08:00
heap_profiler_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
host_callback_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
host_callback_to_tf_test.py
[callback] Add a flag to implement host_callback in terms of io_callback.
2024-04-05 08:51:30 +01:00
image_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
infeed_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
jax_jit_test.py
MAINT Migrate remaining internal/test modules to use state objects
2023-10-12 17:32:15 +01:00
jax_to_ir_test.py
Added int4 and uint4 to dtype-specific tests
2024-04-18 15:20:20 +01:00
jaxpr_effects_test.py
Unify token lowering in JAX to always use stablehlo.token.
2024-04-26 13:18:29 -07:00
jaxpr_util_test.py
MAINT Migrate remaining internal/test modules to use state objects
2023-10-12 17:32:15 +01:00
jet_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
key_reuse_test.py
[key reuse] handle reuse of closed-over constants
2024-04-11 15:39:45 -07:00
lax_autodiff_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_control_flow_test.py
Add specialize on jax.jit so that we can delete the duplicate code in jax.make_jaxpr.
2024-04-27 18:58:16 -07:00
lax_metal_test.py
CI: run the metal CI nightly & on PRs touching the config file
2024-04-11 11:18:27 -07:00
lax_numpy_einsum_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_numpy_indexing_test.py
jax.numpy: implement scalar boolean indexing
2024-02-09 11:00:00 -08:00
lax_numpy_operators_test.py
Fixed hypot bug on nan/inf pairings, began deprecation of non-real values
2024-04-15 17:56:16 +00:00
lax_numpy_reducers_test.py
Add support for copy kwarg in astype to match Array API
2024-04-22 16:25:37 +00:00
lax_numpy_test.py
DOC: improve docs of transpose & matrix_transpose
2024-04-29 08:17:50 -07:00
lax_numpy_ufuncs_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_numpy_vectorize_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_scipy_sparse_test.py
MAINT Migrate remaining internal/test modules to use state objects
2023-10-12 17:32:15 +01:00
lax_scipy_special_functions_test.py
Add jax.scipy.special.gammasgn
2024-04-18 16:14:55 -07: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
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_test.py
Workaround mpmath 1.3 bugs in tan and tanh evaluation at infinities
2024-04-10 18:26:07 +03:00
lax_vmap_op_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_vmap_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
layout_test.py
Add layout support to make_array_from_callback.
2024-04-15 12:38:34 -07:00
lazy_loader_test.py
lazy_loader_module: Move to new internal_test_util directory.
2023-03-10 10:29:33 -08:00
linalg_test.py
Refactor QDWH to be more efficient when run batched under vmap.
2024-04-25 11:48:21 -07:00
lobpcg_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
logging_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
memories_test.py
[XLA:TPU] Support scan loops for parameter input and output streaming in host offloading.
2024-04-21 20:03:11 -07:00
mesh_utils_test.py
Adding expectations to e2e test with expectations for mesh creation allowing splitting physical axes.
2024-03-21 18:58:06 -07:00
metadata_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
mock_gpu_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07: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
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
multibackend_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
multiprocess_gpu_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
name_stack_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
nn_test.py
softmax: deprecate initial argument & always set to -inf internally
2024-04-10 10:23:21 -07:00
ode_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
optimizers_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
package_structure_test.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
pgle_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
pickle_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
pjit_test.py
Add specialize on jax.jit so that we can delete the duplicate code in jax.make_jaxpr.
2024-04-27 18:58:16 -07:00
pmap_test.py
Guard test that requires two devices with device_count() check.
2024-04-01 12:32:54 -07:00
polynomial_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
profiler_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
python_callback_test.py
relax a side-effects test that was erroneously checking against a canonical order
2024-04-15 17:13:35 -07:00
pytorch_interoperability_test.py
Enable some more C API tests.
2024-04-22 09:38:59 -07:00
qdwh_test.py
Refactor QDWH to be more efficient when run batched under vmap.
2024-04-25 11:48:21 -07:00
random_lax_test.py
remove test of deprecated jax.random.shuffle API
2024-04-10 08:19:12 -07:00
random_test.py
Finalize the deprecation of the arr.device() method
2024-04-08 19:04:15 -07:00
scipy_fft_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_interpolate_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_ndimage_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_optimize_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_signal_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_spatial_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_stats_test.py
Merge pull request #20885 from rajasekharporeddy:test_branch4
2024-04-23 13:29:40 -07:00
shape_poly_test.py
[key reuse] rename flag to jax_debug_key_reuse
2024-03-22 05:37:30 -07:00
shard_alike_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
shard_map_test.py
Support auto in shard_map.
2024-04-22 14:29:35 -07:00
source_info_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
sparse_bcoo_bcsr_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
sparse_nm_test.py
Use correct kWidth in sparse dots with int8 input (on Ampere)
2024-04-26 04:53:25 -07:00
sparse_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
sparsify_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
stack_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
state_test.py
Add a zeros rule for mutable arrays and test it using a custom vjp.
2024-04-12 15:22:07 -04:00
stax_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
svd_test.py
Fix bug in rank-deficient fix-up code: Do not zero out the corresponding column of u_out if a diagonal entry of r is exactly zero.
2024-04-18 09:20:48 -07:00
transfer_guard_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
tree_util_test.py
Add builtin cc dataclass pytree node for performance.
2024-04-23 14:14:49 -07:00
typing_test.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
util_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
version_test.py
Add more extensive tests for version strings
2023-09-08 11:33:49 -07:00
x64_context_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
xla_bridge_test.py
Change xla_bridge_test to expect a bytes FDO profile instead of a string in future jaxlib versions.
2024-03-06 13:50:36 -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
xmap_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00