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 most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

api_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

api_util_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

array_interoperability_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

batching_test.py

Clarify the error message for out-of-bounds in_axes in pmap and vmap

2021-07-14 12:11:06 +00:00

callback_test.py

Add support for custom derivatives in jax.experimental.callback

2021-03-22 12:57:33 -07:00

compilation_cache_test.py

Merge pull request #7331 from colemanliyah:fix_teardown

2021-07-19 16:07:34 -07:00

core_test.py

Move jax.abstract_arrays to jax._src.abstract_arrays.

2021-06-02 06:25:22 -07:00

custom_object_test.py

Generalize constant handlers for multi-buffer objects

2021-05-06 09:44:51 -07:00

debug_nans_test.py

jnp.sinc: fix NaNs at x=0

2021-06-10 09:14:07 -07:00

djax_test.py

support implicit broadcasting in transpose rules

2021-04-16 12:51:11 -07:00

dtypes_test.py

Change jnp scalar types to consider numpy scalars as instances.

2021-05-12 20:31:49 -04:00

errors_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

fft_test.py

Add support for padding and cropping to fft

2021-04-17 08:38:24 +02:00

file_system_cache_test.py

implementing thread saftey check

2021-06-28 23:23:49 +00:00

generated_fun_test.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

host_callback_test.py

[host_callback] Add an optional mechanism using HLO CustomCall instead of outfeed.

2021-07-02 23:17:56 -07:00

host_callback_to_tf_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

image_test.py

Add test to ensure gradient is finite

2021-06-01 15:28:25 -06:00

infeed_test.py

Move jax.config to jax._src.config.

2021-04-19 08:53:12 -07:00

jax_jit_test.py

Move jax.abstract_arrays to jax._src.abstract_arrays.

2021-06-02 06:25:22 -07:00

jax_to_hlo_test.py

Small cleanups to dependency structure.

2021-01-20 12:43:28 -08:00

jaxpr_util_test.py

Make jax_enable_x64 a thread-local value.

2021-02-04 09:48:22 -08:00

jet_test.py

Adjust JetTest.test_scatter_add tolerance for TPU

2021-03-31 17:25:27 +00:00

lax_autodiff_test.py

Copybara import of the project:

2021-05-12 02:30:15 -07:00

lax_control_flow_test.py

Add auxiliary data support to lax.custom_root

2021-06-17 19:25:46 +02:00

lax_numpy_einsum_test.py

Try to avoid transposes in jnp.einsum by considering both argument orders to dot_general.

2021-03-04 10:00:08 -05:00

lax_numpy_indexing_test.py

Add x.at[idx].get().

2021-07-07 08:51:09 -04:00

lax_numpy_test.py

lax_numpy_test: disable implicit rank promotion by default

2021-07-13 11:38:21 -07:00

lax_numpy_vectorize_test.py

Mark arguments to jax.jit() other than the function as keyword-only.

2021-04-08 10:32:35 -04:00

lax_scipy_sparse_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

lax_scipy_test.py

Remove additional info return value from jax.scipy.linalg.polar().

2021-07-20 13:13:31 -04:00

lax_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

lax_vmap_test.py

Remove stale limitations

2021-07-08 09:42:26 -07:00

linalg_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

loops_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

masking_test.py

Merge pull request #6375 from gnecula:mask_clean

2021-04-12 05:50:19 -07:00

metadata_test.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00

multi_device_test.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00

multibackend_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

nn_test.py

Remove stale limitations

2021-07-08 09:42:26 -07:00

ode_test.py

Fix scan carry types in gradient of complex ODE (#4130 )

2020-08-24 13:50:44 -07:00

optimizers_test.py

X32 tests: fail on dtype warnings

2020-12-08 13:03:30 -08:00

pjit_test.py

Implement vmap rules for with_sharding_constraint

2021-07-14 06:25:20 -07:00

pmap_test.py

Add axis and tiled options to lax.all_gather.

2021-07-15 04:22:36 -07:00

polynomial_test.py

Skip unimplemented TPU cases in polynomial_test.py

2021-03-04 00:12:55 +00:00

profiler_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

random_test.py

Run random_test with rank_promotion='raise'

2021-07-08 11:20:05 -07:00

scipy_ndimage_test.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

scipy_optimize_test.py

Drop Python 3.6 support.

2021-07-15 14:20:29 -04:00

scipy_signal_test.py

Support complex numbers in jax.scipy.signal.convolve/correlate

2021-06-18 13:07:00 -06:00

scipy_stats_test.py

Fix jax.scipy.stats.gamma.pdf() for x=0.0, a=1.0

2021-07-12 14:43:44 -07:00

sharded_jit_test.py

Enable sharded_jit_test and xmap SPMD tests for GPU.

2021-03-23 10:22:10 -07:00

sparse_test.py

[sparse] Implement several BCOO methods via sparsify

2021-07-02 13:04:35 -07:00

sparsify_test.py

[sparse] support sparse arguments in xla_call

2021-07-13 15:23:14 -07:00

stax_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

test_util_test.py

Merge pull request #7192 from amol-mandhane:broadcasting-decorator

2021-07-07 11:05:34 -07:00

tree_util_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

util_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

x64_context_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

xla_bridge_test.py

Replace host_id with process_index terminology, take 2.

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

xla_interpreter_test.py

Update minimum jaxlib version to 0.1.69.

2021-07-15 17:00:13 -04:00

xmap_test.py

Add an option to strictly enforce sharding implies by named axes

2021-07-19 08:39:27 -07:00