Matthew Johnson 58826507cc [dynamic-shapes] add basic vmap-of-indexing support ...

The main changes here are only indirectly related to gather: we just had to
update some other rules (e.g. for comparison, and squeeze) for a simple
dynamic-batch-shape gather to work.

I also skipped two tests and deleted some old dynamic shape slicing logic
because we want to handle that differently. We didn't have to do that removal
in this PR, but it's just convenient given I'm looking at indexing again.

2022-09-08 17:52:12 -07:00

..

benchmarks

Cleanup: remove unused imports in tests (#3276)

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

filecheck

Lower bessel_i1e primitive to chlo.bessel_i1e.

2022-08-16 12:39:28 -07:00

notebooks

Update colab CPU test notebook to be more robust

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

third_party/scipy

[x64] make scipy_optimize_test compatible with strict dtype promotion

2022-06-21 09:29:08 -07:00

ann_test.py

[JAX] Fix batch logic for approx_min/max_k

2022-04-08 13:50:36 -07:00

api_test.py

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -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

Disabled dlpack bfloat16 test for JAX<->cupy.

2022-08-19 10:08:07 -04:00

array_test.py

Add a test for jnp.array(..., copy=True) for Array.

2022-09-08 11:29:51 -07:00

batching_test.py

roll-forward #11952, take 2

2022-08-22 13:57:31 -07:00

BUILD

Split State effect into Read/Write/Accum effects and tie them to Ref avals

2022-09-08 08:04:13 -07:00

callback_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

checkify_test.py

Checkify: add jvp rule for assert_p.

2022-09-08 05:26:02 -07:00

clear_backends_test.py

Bump the minimum jaxlib version to 0.3.15.

2022-09-08 16:43:46 -04:00

compilation_cache_test.py

[xla:jitrt] Enable serializing GPU executable for JAX

2022-09-08 11:50:31 -07:00

core_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

custom_linear_solve_test.py

add test for custom_linear_solve + new remat

2022-07-28 21:37:23 -07:00

custom_object_test.py

fix custom_object_test for new Array type

2022-08-18 16:33:05 -07:00

custom_root_test.py

Relax tolerances slightly for MKL.

2022-02-25 22:02:55 +00:00

debug_nans_test.py

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -07:00

debugger_test.py

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -07:00

debugging_primitives_test.py

Use .device() to get the device and platform from the device and fix TODO to point to github issue

2022-08-19 13:14:13 -07:00

dtypes_test.py

Add array test counterparts to dtypes_test.

2022-08-18 15:58:39 -07:00

dynamic_api_test.py

[dynamic-shapes] add basic vmap-of-indexing support

2022-09-08 17:52:12 -07:00

errors_test.py

Avoid more direct references to jax._src without imports.

2022-08-24 07:51:28 -07:00

fft_test.py

Add shape checks for lax.fft.

2022-08-13 16:32:52 +00:00

generated_fun_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

gfile_cache_test.py

Using etils(gfile) to support gcs buckets and file system for persistent compilation caching

2022-06-10 00:17:13 +00:00

global_device_array_test.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

heap_profiler_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

host_callback_test.py

delete old remat implementation

2022-08-16 23:16:37 -07:00

host_callback_to_tf_test.py

Internal change

2022-08-18 11:39:07 -07:00

image_test.py

Internal change

2022-08-18 11:39:07 -07:00

infeed_test.py

Avoid top-level aliases of jax.tree_util.*

2022-07-07 11:41:02 -07:00

jax_jit_test.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

jax_to_ir_test.py

Internal change

2022-08-18 11:39:07 -07:00

jaxpr_effects_test.py

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -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

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

lax_autodiff_test.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

lax_control_flow_test.py

Add sound loop invariance detection

2022-09-08 10:42:19 -07:00

lax_numpy_einsum_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

lax_numpy_indexing_test.py

Copybara import of the project:

2022-08-25 15:14:33 -07:00

lax_numpy_test.py

Enable testArrayCopy now that its fixed.

2022-09-08 14:40:06 -07:00

lax_numpy_vectorize_test.py

[x64] make lax_numpy_vectorize_test work with strict dtype promotion

2022-06-16 14:06:24 -07:00

lax_scipy_sparse_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

lax_scipy_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

lax_test.py

Reapply: Use MLIR bytecode when passing IR to backends.

2022-09-08 08:11:16 -07:00

lax_vmap_test.py

Merge pull request #12122 from hawkinsp:fft

2022-08-26 11:32:07 -07:00

linalg_test.py

Bump test tolerance for eigenvalue test.

2022-08-16 14:44:16 -04:00

lobpcg_test.py

Make LOBPCG test plots compatible with bazel.

2022-08-04 20:05:53 -07:00

mesh_utils_test.py

[mesh_utils] Avoid relying on process-tiled device order

2022-06-06 13:30:56 -07:00

metadata_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

multi_device_test.py

Merge pull request #12297 from mattjj:computation-follows-data-prng

2022-09-08 14:57:31 -07:00

multibackend_test.py

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -07:00

multiprocess_gpu_test.py

Add pytest marker inside the test only if pytest is present in the env

2022-09-06 11:45:59 -07:00

name_stack_test.py

Only wrap one level of the name-stack with transform names.

2022-06-10 03:50:26 -07:00

nn_test.py

Update nn_test.py

2022-08-11 15:46:34 +01:00

ode_test.py

[x64] make ode_test pass with strict dtype promotion

2022-06-17 15:21:10 -07:00

optimizers_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

pickle_test.py

support key array pickling

2022-08-31 12:03:53 -07:00

pjit_test.py

Add support to handle arbitrary shardings to KeyArray. Resolve all the TODOs that were created before.

2022-08-31 22:54:06 -07:00

pmap_test.py

Make __iter__ of Array behave like DA when there is a SingleDeviceSharding and like SDA when there is a non-trivial sharding.

2022-09-01 19:54:34 -07:00

polynomial_test.py

Remove _ prefix from functions in jax._src.dtypes.

2022-08-12 12:51:09 +00:00

profiler_test.py

Internal change

2022-08-18 11:39:07 -07:00

python_callback_test.py

Use lowering as impl rule for pure_callback

2022-09-01 15:29:31 -07:00

qdwh_test.py

un-skip polar/qdwh decomp tests skipped on gpu in ad6ce74

2022-08-15 12:31:43 -07:00

random_test.py

Merge pull request #12170 from froystig:just-dtype

2022-08-31 18:36:47 -07: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

[x64] make scipy_ndimage_test compatible with strict dtype promotion

2022-06-16 15:45:15 -07:00

scipy_optimize_test.py

jax.scipy.optimize: fix type inconsistency

2022-06-29 09:47:53 -07:00

scipy_signal_test.py

Remove _ prefix from functions in jax._src.dtypes.

2022-08-12 12:51:09 +00:00

scipy_stats_test.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

sparse_test.py

[sparse] add sparse support for dynamic_slice

2022-09-01 13:42:02 -07:00

sparsify_test.py

[sparse] add sparse support for dynamic_slice

2022-09-01 13:42:02 -07:00

stack_test.py

Make QDWH-eig implementation jit-table.

2022-05-17 18:12:48 -07:00

state_test.py

Split State effect into Read/Write/Accum effects and tie them to Ref avals

2022-09-08 08:04:13 -07:00

stax_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

svd_test.py

[x64] make linalg functions & tests compatible with strict dtype promotion

2022-06-16 10:32:20 -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

Make PyTreeDef pickleable

2022-08-04 07:13:46 -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

feat: validate jit args

2022-05-18 21:54:47 +01:00

xla_bridge_test.py

Avoid more direct references to jax._src without imports.

2022-08-24 07:51:28 -07:00

xla_interpreter_test.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

xmap_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00