rocm_jax

deprecation_module

Add infrastructure for managing deprecations.

2023-02-09 05:48:40 -08:00

filecheck

[CHLO] Add erf_inv and lowering to mhlo

2023-03-01 02:52:52 -08:00

lazy_loader_module

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

notebooks

Make colab_gpu.ipynb compatible with newer JAX versions

2023-02-13 15:56:58 -08:00

third_party/scipy

[x64] more type safety in scipy.optimize.line_search

2022-12-01 14:04:39 -08:00

ann_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

api_test.py

attach debug info to jaxpr, pass to mlir/mhlo

2023-03-02 17:23:58 -08:00

api_util_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

array_interoperability_test.py

Catch ImportError when importing tf instead of a broad exception catch. If not, this leads to weird errors in the other tests down the line.

2023-02-16 12:03:58 -08:00

array_test.py

[XLA:Python] Add buffer protocol support to jax.Array.

2023-02-28 17:35:40 -08:00

batching_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

BUILD

lax_scipy_test: Revert split into three targets.

2023-02-24 16:56:40 -08:00

checkify_test.py

Rename in_axis_resources and out_axis_resources with in_shardings and out_shardings. This is just a simple name replacement. It does not change any of the current pjit semantics and doesn't break any code.

2023-02-18 10:00:36 -08:00

clear_backends_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

compilation_cache_test.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

core_test.py

remove several more symbols from jax.core

2023-02-17 13:07:00 -08:00

custom_linear_solve_test.py

Fix issue in check_tree, so that custom_linear_solve supports

2022-11-07 21:29:43 +01:00

custom_object_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

custom_root_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

debug_nans_test.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

debugger_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

debugging_primitives_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

deprecation_test.py

Add infrastructure for managing deprecations.

2023-02-09 05:48:40 -08:00

dtypes_test.py

Replace uses of jnp.ndarray with jax.Array inside JAX.

2023-02-15 14:53:00 -08:00

dynamic_api_test.py

remove core.bint

2023-02-15 14:28:29 -08:00

errors_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

experimental_rnn_test.py

Opt into higher matmul precision for A100 and TPU tests.

2023-02-14 12:03:12 -08:00

fft_test.py

Remove placeholder functions for unimplemented NumPy functions.

2023-02-02 13:00:18 -05:00

for_loop_test.py

[x64] minor test-only updates for better type safety

2022-11-30 15:18:40 -08:00

generated_fun_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

gfile_cache_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

global_device_array_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

heap_profiler_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

host_callback_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

host_callback_to_tf_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

image_test.py

Add changes accidentally omitted from

2022-10-10 19:11:58 +00:00

infeed_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

jax_jit_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

jax_to_ir_test.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

jaxpr_effects_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

jaxpr_util_test.py

Use jnp.arange to break the pjit cache (when jit and pjit are merged) because pytest runs tests non-hermetically.

2023-02-08 10:17:37 -08:00

jet_test.py

relax tanh test tols for upcoming xla change

2022-12-20 21:06:09 -08:00

lax_autodiff_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

lax_control_flow_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

lax_numpy_einsum_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax_numpy_indexing_test.py

Roll-back https://github.com/google/jax/pull/14144 due to downstream test failures

2023-01-25 12:15:36 -08:00

lax_numpy_operators_test.py

Fix test failures under NumPy 1.24.

2023-01-18 14:25:33 -08:00

lax_numpy_reducers_test.py

jax.numpy reductions: require initial to be a scalar

2023-02-14 15:36:18 -08:00

lax_numpy_test.py

Don't create invalid bools in lax_numpy_test/testView.

2023-02-28 14:03:08 -08:00

lax_numpy_vectorize_test.py

Port more tests from jtu.cases_from_list to jtu.sample_product.

2022-10-11 21:06:08 +00:00

lax_scipy_sparse_test.py

add "linear solve batching via jacrev" test from github.com/google/jax/issues/14249

2023-02-01 20:01:53 -08:00

lax_scipy_test.py

lax_scipy_test: Revert split into three targets.

2023-02-24 16:56:40 -08:00

lax_test.py

Add sharding to the signature of shard_args and update

2023-03-02 13:29:03 -08:00

lax_vmap_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

lazy_loader_test.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

linalg_test.py

[ROCm] Re-enable some linalg and sparse tests

2023-02-07 22:05:14 +00:00

lobpcg_test.py

Replace uses of jnp.ndarray with jax.Array inside JAX.

2023-02-15 14:53:00 -08:00

mesh_utils_test.py

[JAX] Deprecate jax.experimental.maps.Mesh.

2023-02-15 09:15:50 -08:00

metadata_test.py

Remove uses of jax.xla_computation from metadata_test.py

2023-02-15 17:26:21 -08:00

monitoring_test.py

Add JAX monitoring library that instruments code via events.

2022-11-15 12:41:41 -08:00

multi_device_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

multibackend_test.py

pjit allows nesting of pjits where the outer backend is None while the inner backend is something other than device_under_test(). This is because the inner backend will take priority.

2023-01-17 16:39:45 -08:00

multiprocess_gpu_test.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

name_stack_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

nn_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

ode_test.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

2022-10-28 16:05:25 -07:00

optimizers_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

pickle_test.py

Replace op_sharding_sharding with gspmd_sharding. This is purely an internal change.

2023-02-17 17:53:13 -08:00

pjit_test.py

Internal change

2023-03-04 13:24:11 +00:00

pmap_test.py

Add sharding to the signature of shard_args and update

2023-03-02 13:29:03 -08:00

polynomial_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

profiler_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

python_callback_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

pytorch_interoperability_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

qdwh_test.py

Replace uses of jnp.ndarray with jax.Array inside JAX.

2023-02-15 14:53:00 -08:00

random_test.py

add random.chisquare and random.f

2023-03-01 15:03:50 +08:00

remote_transfer_test.py

Migrate pytestmark usage to new @jtu.pytest_mark_if_available decorator.

2023-01-12 22:44:39 +00:00

scipy_fft_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

scipy_interpolate_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

scipy_ndimage_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

scipy_optimize_test.py

Port more tests from jtu.cases_from_list to jtu.sample_product.

2022-10-11 21:06:08 +00:00

scipy_signal_test.py

Rename test flag --num_generated_cases to --jax_num_generated_cases.

2022-11-09 18:58:05 -08:00

scipy_stats_test.py

Fixed loc and scale parameters for logistic distribution. CDF and SF have been added for several distributions, including cauchy, gamma, logistic, chi2 and beta. ISF and PPF have also been added for cauchy and logistic.

2023-03-07 07:56:47 +01:00

shard_map_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

sparse_test.py

[sparse] more robust correction for out-of-bound indices in BCOO

2023-03-02 14:26:25 -08:00

sparsify_test.py

[sparse] implement bcsr_concatenate

2023-02-15 14:10:56 -08:00

stack_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

state_test.py

Refactor Ref abstract type to contain other AbstractValues

2023-02-23 17:02:40 -08:00

stax_test.py

[x64] more type safety in stax_test.py

2022-12-02 10:02:25 -08:00

svd_test.py

[x64] make tests more type-safe

2022-12-02 13:21:35 -08:00

transfer_guard_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

tree_util_test.py

Internal change

2023-03-04 13:24:11 +00:00

typing_test.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

util_test.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

version_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

x64_context_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

xla_bridge_test.py

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

xla_interpreter_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

xmap_test.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00