rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

2020-06-01 11:49:35 -07:00

notebooks

Update colab CPU test notebook to be more robust

2021-07-13 09:30:56 -07:00

third_party/scipy

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

ann_test.py

[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.

2021-10-28 16:49:37 -07:00

api_test.py

Merge pull request #8403 from shoyer:empty-map-args-error

2021-11-08 19:39:04 -08: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

add internal vmappable interface (part 1)

2021-11-04 15:01:54 -07:00

callback_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

compilation_cache_test.py

Include XLA_FLAGS in persistent compilation cache key.

2021-10-06 14:11:40 -07:00

core_test.py

Change jax.core.DropVar to be a non-singleton.

2021-10-18 15:02:54 -07:00

custom_object_test.py

Port remaining translation rules inside JAX to new style.

2021-10-19 09:48:37 -07:00

debug_nans_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

djax_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

dtypes_test.py

jnp.result_type: respect default types

2021-11-01 11:44:14 -07:00

errors_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

fft_test.py

add fft normalisation

2021-10-20 22:15:35 +01: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

gsda_test.py

GSDA callback implementation for from_batched_callback and from_batched_callback_with_devices.

2021-11-05 00:15:40 -07: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

Move xla_bridge.constant to jax.interpreter.xla.pyval_to_ir_constant.

2021-10-19 08:40:51 -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

[JAX] Canonicalize shapes in jax.image.resize().

2021-09-30 12:38:39 -07:00

infeed_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jax_jit_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jax_to_hlo_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jaxpr_util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jet_test.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

lax_autodiff_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

lax_control_flow_test.py

fix #8649

2021-11-08 14:52:00 -08:00

lax_numpy_einsum_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

lax_numpy_indexing_test.py

jax.numpy: fix boolean indexing with Ellipsis

2021-11-02 09:15:08 -07:00

lax_numpy_test.py

jnp.result_type: respect default types

2021-11-01 11:44:14 -07:00

lax_numpy_vectorize_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

lax_scipy_sparse_test.py

Reenable a test that was disabled due to an (apparently fixed) LLVM bug.

2021-10-16 10:34:36 -07:00

lax_scipy_test.py

Merge pull request #8291 from jakevdp:lpmn-tests

2021-10-20 11:35:50 -07:00

lax_test.py

Merge pull request #8043 from hawkinsp:iter

2021-11-01 07:41:40 -07:00

lax_vmap_test.py

Implement lax.reduce() batching rule case where batch dimensions differ between operands.

2021-11-03 09:42:01 -04:00

linalg_test.py

Implement hermitian=... option on jax.numpy.linalg.svd.

2021-11-01 09:55:30 -04:00

loops_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

masking_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

metadata_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

multi_device_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

multibackend_test.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

nn_test.py

fix whitespace

2021-10-27 22:06:06 -07:00

ode_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

optimizers_test.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

pickle_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

pjit_test.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

pmap_test.py

add more grad-of-jit/pmap caching tests

2021-10-13 11:06:17 -07:00

polynomial_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

profiler_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

qdwh_test.py

[JAX] Added jax.lax.linalg.qdwh.

2021-10-29 14:45:06 -07:00

random_test.py

jax.random.permutation: add independent keyword

2021-11-02 11:39:41 -07:00

scipy_fft_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_ndimage_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_optimize_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_signal_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_stats_test.py

stats.multivariate_normal: support broadcasted inputs

2021-10-19 16:58:36 -07:00

sharded_jit_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

sparse_test.py

[sparse] add full support for BCOO matmul operator

2021-11-04 10:54:23 -07:00

sparsify_test.py

[sparse]: add tracer-based implementation of sparsify

2021-11-04 13:02:46 -07:00

stax_test.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

tree_util_test.py

Fixes for Python 3.10.

2021-10-05 15:25:28 -04:00

util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

x64_context_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

xla_bridge_test.py

Remove xla_bridge.make_computation_builder().

2021-10-18 13:20:34 -04:00

xla_interpreter_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

xmap_test.py

Use shaped_abstractify to allow abstract arguments in .lower methods

2021-09-27 05:32:12 -07:00