mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
jax authors 11da3df238 Merge pull request #21096 from gspschmid:gschmid/sourcemaps 
PiperOrigin-RevId: 631769572
2024-05-08 05:44:08 -07:00
..
filecheck
tests: access tree utilities via jax.tree.*
2024-02-26 14:17:18 -08:00
mosaic
[jax:mosaic-gpu] FragmentedArray can do tiled load.
2024-05-07 18:13:55 -07:00
notebooks
Deprecate arr.device_buffer and arr.device_buffers
2023-12-06 10:20:29 -08:00
pallas
[Pallas] Fix some typos.
2024-05-07 16:52:38 -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
Refactor array_api namespace, relying more directly on jax.numpy
2024-05-02 18:17:45 +00:00
array_interoperability_test.py
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -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
Merge pull request #21096 from gspschmid:gschmid/sourcemaps
2024-05-08 05:44:08 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07:00
core_test.py
Allow replacing jaxpr debug_info with None.
2024-04-30 08:31:39 -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
[export] Add backwards compatibility test for Pallas call on GPUs.
2024-05-02 05:38:33 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07: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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07: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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -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
jet_test: use float32 matmul precision
2024-05-02 14:48:22 +02: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
metal_plugin ci with jaxlib nightly
2024-05-02 18:40:01 -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
Handle ellipsis ... in _attempt_rewriting_take_via_slice.
2024-05-07 10:30:08 -07: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
Refactored common upcast for integral-type accumulators
2024-05-06 15:13:10 +00:00
lax_numpy_test.py
Add optional size argument to jnp.compress & jnp.extract.
2024-05-07 08:47:34 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07: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
jnp.linalg tensorinv & tensorsolve: improve implementation & docs
2024-05-06 11:08:36 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -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
mutable_array_test.py
Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file.
2024-05-02 14:36:16 -04: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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07:00
pmap_test.py
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07: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
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
python_callback_test.py
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -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
random_lax_test: fix kstest for newer NumPy
2024-05-06 09:20:07 -07:00
random_test.py
reintroduce the Threefry GPU kernel lowering, under a flag
2024-05-01 10:33:31 -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
DOC: Improve remaining jax.scipy docstrings
2024-05-02 10:29:00 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07:00
shard_map_test.py
[shard_map] better fix for spmd_axis_name issues with shmap residuals
2024-05-04 01:31:15 +00:00
source_info_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07: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
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 discharge rules for scan with mutable arrays. Move mutable array tests to separate file.
2024-05-02 14:36:16 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -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
Bump minimum jaxlib version to 0.4.27
2024-05-07 16:07:59 -07: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