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
Yash Katariya 6072d5993e Any devices passed to jax.sharding.Mesh are required to be hashable. 
This is true for mock devices or user specific devices and jax.devices() too.

Fix the tests so that the mock devices are hashable.

PiperOrigin-RevId: 561103167
2023-08-29 12:20:54 -07:00
..
filecheck
Make mlir.custom_call() more general and expose it as jax.interpreters.mlir.custom_call().
2023-08-29 08:50:07 -07:00
notebooks
Delete remote TPU support.
2023-03-24 12:33:33 -07:00
pallas
Load Python rules from rules_python
2023-08-24 10:22:57 -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
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
ann_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
aot_test.py
Add internal jaxlib function for fetching the topology from
2023-08-01 14:54:08 -07:00
api_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
api_util_test.py
jax.jit now works correctly if both donate_argnums and donate_argnames are specified.
2023-07-14 14:28:16 -07:00
array_interoperability_test.py
Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54
2023-08-01 08:53:28 -07:00
array_test.py
Don't drop out of C++ fast path if mesh pointers are not equal.
2023-08-28 15:04:05 -07:00
batching_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
BUILD
create jax.extend.random as a copy of jax.prng
2023-08-24 14:41:56 -07:00
cache_key_test.py
Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.
2023-08-15 06:39:46 -07:00
checkify_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
clear_backends_test.py
Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54
2023-08-01 08:53:28 -07:00
compilation_cache_test.py
Create a context manager to register/unregister event duration listeners.
2023-08-22 13:35:45 -07:00
core_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
custom_linear_solve_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
custom_object_test.py
Move functions out of xla.py closer to their users.
2023-08-08 14:40:42 -07:00
custom_root_test.py
Relax test tolerances.
2023-06-23 09:22:11 -07:00
debug_nans_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
debugger_test.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
debugging_primitives_test.py
Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332
2023-06-01 09:42:55 -07:00
deprecation_test.py
deprecation_module: Move to new internal_test_util directory.
2023-02-17 10:55:04 -08:00
dtypes_test.py
Remove dtypes_test.py:testFinfo
2023-08-24 12:13:54 -07:00
dynamic_api_test.py
fix pow_p jvp rule at x=0. y=0
2023-07-28 17:14:47 -07:00
errors_test.py
Tracebacks no longer have JAX-internal frames prepended by default
2023-08-03 11:38:38 -07:00
experimental_rnn_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
extend_test.py
add jax.extend.random.wrap_key_data
2023-08-26 11:39:25 -07:00
fft_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
for_loop_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
generated_fun_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -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
heap_profiler_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
host_callback_test.py
[callback] Add workaround for TPU host_callback not supporting empty arrays.
2023-08-22 03:47:18 -07:00
host_callback_to_tf_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
image_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
infeed_test.py
Remove 'pjrt_c_api_unimplemented' pytest mark.
2023-03-24 23:14:54 +00:00
jax_jit_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
jax_to_ir_test.py
Delete jax_experimental_name_stack flag
2022-11-10 11:59:50 -08:00
jaxpr_effects_test.py
Fix scan DCE rule to update effects
2023-05-16 23:17:02 +00:00
jaxpr_util_test.py
Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target.
2023-08-08 10:09:09 -07:00
jet_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
lax_autodiff_test.py
lax.pow: fix shape mismatch failure in jvp rule
2023-08-25 10:05:55 -07:00
lax_control_flow_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
lax_numpy_einsum_test.py
jnp.einsum: lower to mixed-precision dot_general when possible.
2023-08-24 10:31:39 -07:00
lax_numpy_indexing_test.py
ENH: Update numpy exceptions imports
2023-08-07 19:08:41 +02:00
lax_numpy_operators_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
lax_numpy_reducers_test.py
ENH: Update numpy exceptions imports
2023-08-07 19:08:41 +02:00
lax_numpy_test.py
Remove trivial execution from jax since it leads to 100x slower dispatch time.
2023-08-25 10:59:48 -07:00
lax_numpy_ufuncs_test.py
jnp.ufunc: add fast paths for add/prod reductions
2023-08-28 08:30:23 -07:00
lax_numpy_vectorize_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
lax_scipy_sparse_test.py
Relax test tolerances.
2023-06-23 09:22:11 -07:00
lax_scipy_special_functions_test.py
Fix tolerance on bernoulli test
2023-08-23 16:59:27 -07:00
lax_scipy_spectral_dac_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
lax_scipy_test.py
Deprecate jax.numpy.trapz.
2023-08-25 09:04:13 -06:00
lax_test.py
Remove trivial execution from jax since it leads to 100x slower dispatch time.
2023-08-25 10:59:48 -07:00
lax_vmap_op_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
lax_vmap_test.py
Fix scatter batching rule for scatter_apply
2023-07-10 16:42:45 -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
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
lobpcg_test.py
Remove matplotlib from the test requirements.
2023-06-14 09:02:49 -04:00
logging_test.py
Add jax_debug_log_modules config option.
2023-07-28 18:11:12 +00:00
mesh_utils_test.py
Any devices passed to jax.sharding.Mesh are required to be hashable.
2023-08-29 12:20:54 -07:00
metadata_test.py
Change metadata_test to tolerate paths with backslashes.
2023-06-12 21:57:41 -04:00
monitoring_test.py
Add an unregister event listener function in JAX monitoring.
2023-08-18 20:00:03 -07:00
mosaic_test.py
Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic
2023-07-20 18:28:51 -07:00
multi_device_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
multibackend_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
multiprocess_gpu_test.py
Replace references to DeviceArray with Array.
2023-08-18 17:46:00 -04:00
name_stack_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
nn_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
ode_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
optimizers_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
package_structure_test.py
Apply pyupgrade --py39-plus.
2023-07-21 14:49:44 -04:00
pgle_test.py
Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54
2023-08-01 08:53:28 -07:00
pickle_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
pjit_test.py
Remove trivial execution from jax since it leads to 100x slower dispatch time.
2023-08-25 10:59:48 -07:00
pmap_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
polynomial_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
profiler_test.py
Bump ProfilerTest.test_remote_profiler to run computations for 10 seconds.
2023-07-27 23:54:42 +00:00
python_callback_test.py
Add an optional sharding argument to pure_callback and io_callback
2023-08-25 14:57:35 -07:00
pytorch_interoperability_test.py
[custom-jvp/vjp] for symbolic zeros, ensure rules can be run more than once
2023-08-21 15:28:43 -07:00
qdwh_test.py
Fix test failures in JAX under NumPy 1.25.0rc1.
2023-06-09 14:10:35 -07:00
random_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
scipy_fft_test.py
Disable testiDctn just like testiDct
2023-07-31 13:58:46 -07:00
scipy_interpolate_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
scipy_ndimage_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
scipy_optimize_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
scipy_signal_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
scipy_spatial_test.py
Fix problematic rotation tests
2023-06-20 05:27:20 -05:00
scipy_stats_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
shard_map_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
source_info_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
sparse_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
sparsify_test.py
Previously, using sparse.todense on a BCSR matrix with sparse.sparsify would raise NotImplementedError: sparse rule for todense is not implemented. By adding the sparse rule, it will resolve this issue.
2023-07-26 13:01:02 -07:00
stack_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
state_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
stax_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
svd_test.py
[ROCm]: Re-enable SvdOnTinyElement test
2023-06-23 17:50:00 +00:00
transfer_guard_test.py
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
tree_util_test.py
Add register_static decorator to tree_util to facilitate creating leafless classes.
2023-08-21 16:54:29 -07:00
typing_test.py
Remove references to jax.config.jax_array, which is always True at head.
2023-03-15 17:09:11 -07:00
util_test.py
Remove obsolete jaxlib version checks
2023-07-12 11:53:55 -07:00
version_test.py
Check that jaxlib's nightly date is not greater than jax's nightly date in check_jaxlib_version
2023-08-18 10:54:52 -07:00
x64_context_test.py
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00
xla_bridge_test.py
Ensure that CompileOptions serializes deterministically.
2023-08-23 11:34:21 -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
Tests: set jax_legacy_prng_key='error'
2023-08-28 10:56:09 -07:00