mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
Peter Hawkins a52f07a21b Add an optional mode= argument to jnp.take_along_axis. 
This allows users of jnp.take_along_axis to override the out-of-bounds indexing behavior.
Default to "clip", which for the forward computation is identical to the current behavior. In a future change, we will change this to "fill".
2022-04-19 16:07:00 -04:00
..
benchmarks
Cleanup: remove unused imports in tests (#3276)
2020-06-01 11:49:35 -07:00
filecheck
[MHLO] Switch call_tf to use an MHLO lowering.
2022-04-19 07:17:22 -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
[MHLO] Remove most XLA translation rules.
2022-04-18 08:28:35 -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
[XLA] Call translation rule directly in xla.primitive_subcomputation.
2022-04-18 18:48:27 -07:00
custom_root_test.py
Relax tolerances slightly for MKL.
2022-02-25 22:02:55 +00:00
debug_nans_test.py
debug_nans: don't return results of successfully running de-optimized function
2022-04-12 14:40:19 -07: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
Increase minimum jaxlib version to 0.3.7.
2022-04-18 08:09:50 -04: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
[MHLO] Change jax.xla_computation() to use MHLO lowering internally.
2022-04-13 06:28:38 -07: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
Fix name stack bugs
2022-04-19 11:14:41 -07:00
jaxpr_effects_test.py
Add tests for higher order primitives
2022-04-12 18:12:44 -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
better errors: check for callability of lax.control_flow arguments
2022-04-13 10:39:01 -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
Do not generate trivial gathers when indexing entire axis
2022-04-10 01:06:39 +01:00
lax_numpy_test.py
Add an optional mode= argument to jnp.take_along_axis.
2022-04-19 16:07:00 -04: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
Prevent negative output shapes in shape inference for reduce_window.
2022-04-19 12:36:31 -04: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
masking_test.py
Testing: avoid hard-coding random seeds
2021-12-10 10:32:09 -08:00
mesh_utils_test.py
[mesh_utils] Support creating device meshes for hybrid networks
2022-04-18 11:00:30 -07:00
metadata_test.py
Fix name stack bugs
2022-04-19 11:14:41 -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
[MHLO] Change jax.xla_computation() to use MHLO lowering internally.
2022-04-13 06:28:38 -07: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
Drop dead code now that the minimum jaxlib version is 0.3.2.
2022-04-13 13:34:00 -04: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
Fix PRNGKeyArray.broadcast_to with scalar shape
2022-04-16 00:27:30 +08:00
remote_transfer_test.py
Increase minimum jaxlib version to 0.3.7.
2022-04-18 08:09:50 -04: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
Fix#10219
2022-04-13 04:04:11 +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
[MHLO] Remove most XLA translation rules.
2022-04-18 08:28:35 -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
[linalg] Adds compute_uv to TPU SVD.
2022-04-19 11:28:43 -07:00
transfer_guard_test.py
Bump minimum jaxlib version to 0.3.2 and remove transfer guard compatibility code
2022-04-11 15:33:27 +00: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
[MHLO] Change jax.xla_computation() to use MHLO lowering internally.
2022-04-13 06:28:38 -07:00