mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 21:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
Peter Hawkins 9455254b9f [MHLO] Add a direct MHLO lowering for pjit_p, which lacked one. 
This is a second attempt at this change. In this version, check for and report an error on jit(pjit(...)), which was the root cause of the failure that led to the previous version being reverted.

PiperOrigin-RevId: 441214076
2022-04-12 10:30:52 -07:00
..
benchmarks
Cleanup: remove unused imports in tests (#3276)
2020-06-01 11:49:35 -07:00
filecheck
Remove the jax_enable_mlir flag. MLIR is now the only supported code path.
2022-04-04 08:40:09 -07:00
notebooks
Update colab CPU test notebook to be more robust
2021-07-13 09:30:56 -07:00
third_party/scipy
Fix test failure in line search test due to scipy 1.8 release.
2022-02-07 14:49:54 +00:00
ann_test.py
[JAX] Fix batch logic for approx_min/max_k
2022-04-08 13:50:36 -07:00
api_test.py
prototyping dynamic shapes
2022-04-11 22:10:47 -07:00
api_util_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
array_interoperability_test.py
update test of dlpack error message
2021-10-07 19:14:13 -07:00
batching_test.py
vmap: preserve weak_type in batching tracer
2022-03-30 11:06:56 -07:00
callback_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
checkify_test.py
Checkify: add axis and axis size to OOB error message.
2022-03-31 15:16:35 +01:00
compilation_cache_test.py
Run compilation_cache_test.py using TFRT
2022-02-04 10:17:59 -08:00
core_test.py
prototyping dynamic shapes
2022-04-11 22:10:47 -07:00
custom_linear_solve_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
custom_object_test.py
[MHLO] Add direct MHLO lowerings for sparse primitives.
2022-04-08 08:43:57 -07:00
custom_root_test.py
Relax tolerances slightly for MKL.
2022-02-25 22:02:55 +00:00
debug_nans_test.py
Deprecate maps.mesh and replace it with maps.Mesh.
2022-02-23 10:47:06 -08:00
dtypes_test.py
remove _convert_element_type from public jax.lax module
2022-03-09 18:46:38 -08:00
errors_test.py
Increase minimum jaxlib version to 0.3.0.
2022-03-04 10:33:03 -05:00
fft_test.py
Skip test outside x64
2022-04-04 16:00:18 -07:00
file_system_cache_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
generated_fun_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
global_device_array_test.py
Make mesh_axes on GDA strict by only allowing PartitionSpecs to be consistent with pjit.
2022-03-07 08:59:23 -08:00
heap_profiler_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
host_callback_test.py
Avoid retracing when a host_callback.call is called multiple times with the same function.
2022-04-01 14:41:14 -04:00
host_callback_to_tf_test.py
Move jax.test_util to jax._src.test_util.
2021-09-24 07:02:49 -07:00
image_test.py
Make negative dimension return consistent results for image.scale_and_translate
2022-02-24 22:34:00 -05:00
infeed_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
jax_jit_test.py
jax_jit_test: inherit from JaxTestCase
2022-01-24 15:32:59 -08:00
jax_to_ir_test.py
Add build time support for AOT compilation to TF graphs.
2021-12-20 06:18:39 -08:00
jaxpr_effects_test.py
Adds simple effect types to jaxprs
2022-04-11 11:50:41 -07:00
jaxpr_util_test.py
Increase minimum jaxlib version to 0.3.0.
2022-03-04 10:33:03 -05:00
jet_test.py
Add test and jet-primitive for dynamic_slice
2022-02-08 13:28:41 +01:00
lax_autodiff_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
lax_control_flow_test.py
avoid batching units in cond partial eval
2022-03-22 17:42:38 -07:00
lax_numpy_einsum_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
lax_numpy_indexing_test.py
segment_max: fix identity for boolean dtype
2022-03-15 09:20:20 -07:00
lax_numpy_test.py
Better error messages for jnp.fromiter and jnp.fromfile
2022-03-29 14:30:32 -07:00
lax_numpy_vectorize_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
lax_scipy_sparse_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
lax_scipy_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
lax_test.py
Adds simple effect types to jaxprs
2022-04-11 11:50:41 -07:00
lax_vmap_test.py
fixup from 5415306: remove extraneous lines
2022-03-11 15:19:10 -08:00
linalg_test.py
Add jax.scipy.linalg.rsf2csf
2022-04-06 21:06:23 +01:00
loops_test.py
[x64] make jax.experimental.loops consistent with default dtype
2021-12-08 12:08:49 -08:00
masking_test.py
Testing: avoid hard-coding random seeds
2021-12-10 10:32:09 -08:00
mesh_utils_test.py
Allow a single logical mesh dim to take all devices.
2022-03-16 20:58:59 -07:00
metadata_test.py
[JAX] Fix pylint errors.
2021-10-04 17:54:46 -07:00
multi_device_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
multibackend_test.py
Cleanup internal representation of XLA translation rules.
2021-10-16 07:53:24 -07:00
name_stack_test.py
Add separate mechanism for threading name stacks to the lowering
2022-02-23 09:59:09 -08:00
nn_test.py
Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.
2022-03-23 20:55:22 +00:00
ode_test.py
[XLA:CPU] Relax test tolerances for tests using XLA:CPU.
2022-02-28 09:26:54 -08:00
optimizers_test.py
Testing: avoid hard-coding random seeds
2021-12-10 10:32:09 -08:00
pickle_test.py
Increase minimum jaxlib version to 0.1.74.
2021-11-18 15:06:58 -05:00
pjit_test.py
[MHLO] Add a direct MHLO lowering for pjit_p, which lacked one.
2022-04-12 10:30:52 -07:00
pmap_test.py
Deprecate jax.tree_util.tree_multimap
2022-04-01 14:51:54 -07:00
polynomial_test.py
Testing: avoid hard-coding random seeds
2021-12-10 10:32:09 -08:00
profiler_test.py
Fix keyword argument confusion in jax.profiler.annotate_function decorator.
2021-12-06 09:10:44 -05:00
qdwh_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
random_test.py
Enable tests related to the Gamma distribution for non-default PRNG implementations only when jax_enable_custom_prng is enabled, for consistency with other tests.
2022-04-08 01:08:55 -07:00
scipy_fft_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
scipy_interpolate_test.py
Merge pull request #8625 from Edenhofer:regular_grid_interpolator
2022-01-24 12:02:34 -08:00
scipy_ndimage_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
scipy_optimize_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
scipy_signal_test.py
Add istft to jax.scipy.signal.
2022-04-01 14:28:53 +09:00
scipy_stats_test.py
disable implicit rank promotion in a number of remaining tests
2022-01-28 08:16:30 -08:00
sharded_jit_test.py
JaxTestCase: set jax_numpy_rank_promotion='raise' by default
2022-02-14 09:22:05 -08:00
sparse_test.py
[sparse] Use sorted indices instead of sorted rows only.
2022-04-09 08:33:48 -07:00
sparsify_test.py
[sparse] refactor internal implementation of sparsify transform
2022-03-07 12:48:03 -08:00
stax_test.py
disable implicit rank promotion in a number of remaining tests
2022-01-28 08:16:30 -08:00
svd_test.py
Update JAX to use new math libraries in ROCm-5.0.
2022-03-01 20:02:15 +00:00
transfer_guard_test.py
Implement the JAX transfer guard API
2022-02-14 13:11:49 -08:00
tree_util_test.py
Deprecate jax.tree_util.tree_multimap
2022-04-01 14:51:54 -07:00
util_test.py
Fix race condition for weakref destructor by catching rare exceptions.
2022-04-01 12:04:36 -07:00
version_test.py
Introduce a new jax/jaxlib versioning scheme.
2022-02-07 17:59:42 -05:00
x64_context_test.py
[x64] make x64_context_test more robust
2021-11-22 14:54:30 -08:00
xla_bridge_test.py
Consolidate more XLA-lowering logic between jit, pmap, and xmap.
2021-11-30 14:24:33 -08:00
xla_interpreter_test.py
Split jax.interpreters.xla up into three pieces:
2021-11-22 08:22:43 -08:00
xmap_test.py
Use sharded shape to compute aliasing.
2022-03-10 08:41:57 -08:00