Yash Katariya 958234a9c1 Thread the mesh context manager to the place where we recover out_shardings back from GSPMDShardings. Before if you had a program like this: ...

```
with mesh:
  out = pjit(lambda: 1)()
```

The sharding of `out` was a `GSPMDSharding` which is not ideal. This change fixes that and returns a `NamedSharding` instead.

This is also required for `Shardy` integration.

PiperOrigin-RevId: 658842350

2024-08-02 11:04:48 -07:00

..

filecheck

Integrate StableHLO at openxla/stablehlo@8817ff1d

2024-07-15 10:38:09 -07:00

mosaic

[Mosaic GPU] Skip matmul tests with large clusters

2024-07-26 09:21:13 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

2023-12-06 10:20:29 -08:00

pallas

[jax] Shard pallas_vmap_test

2024-08-02 10:41:22 -07:00

testdata

[PJRT C API] Add parsing PJRT client create options from json file.

2023-03-21 16:57:34 -07:00

third_party/scipy

Fix CI test failure in line_search_test.

2024-05-15 07:38:20 -07:00

ann_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

aot_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

api_test.py

Fix issue with multiple arguments when using custom_vjp with optimize_remat.

2024-08-01 08:44:35 -04:00

api_util_test.py

Replace donation_vector's logic with donation_vector_with_in_tree which is now deleted

2024-04-23 17:38:30 -07:00

array_api_skips.txt

[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api

2024-08-01 11:19:17 -07:00

array_api_test.py

[array api] fix deprecation to support old import pattern

2024-08-01 14:38:59 -07:00

array_interoperability_test.py

Update minimum NumPy version to v1.24.

2024-06-21 15:17:17 -07:00

array_test.py

#sdy Add dataclass that stores the data needed to build a Shardy MLIR sharding attribute.

2024-07-23 12:22:05 -07:00

attrs_test.py

Make attrs work with pytrees

2024-05-28 23:23:51 -04:00

batching_test.py

remove pdot and xeinsum (since xmap is gone)

2024-07-25 21:19:17 +00:00

blocked_sampler_test.py

[Pallas] Implement block-invariant sampling.

2024-06-24 11:20:39 -07:00

BUILD

Un-skip //third_party/py/jax/tests:pytorch_interoperability_test_cpu on ASAN.

2024-07-26 11:10:41 -07:00

cache_key_test.py

Strip device_assignment on GPU platform.

2024-03-19 01:50:39 -07:00

checkify_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

cholesky_update_test.py

Bump minimum jaxlib version to v0.4.30.

2024-06-18 12:35:08 -04:00

clear_backends_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

compilation_cache_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

config_test.py

Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.

2024-07-22 15:09:15 -07:00

core_test.py

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

custom_linear_solve_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

custom_object_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

custom_root_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

debug_nans_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

debugger_test.py

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -07:00

debugging_primitives_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

deprecation_test.py

Fix incorrect usage of assertWarns in tests

2024-07-24 14:16:51 -07:00

device_test.py

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00

dtypes_test.py

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

dynamic_api_test.py

Simplify implementation of _broadcast_to.

2024-07-24 10:57:54 -04:00

errors_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

experimental_rnn_test.py

MAINT Do not import the config object in JAX internals

2023-10-18 10:55:13 +01:00

export_back_compat_test.py

Activate LU Decomposition to XLA's FFI

2024-08-02 02:22:53 -07:00

export_harnesses_multi_platform_test.py

Remove forward compatibility mode for old PRGN custom call on GPU

2024-07-31 08:10:17 -07:00

export_test.py

[export] Fix poly shape check for vjp function with integer valued, polymorphic output.

2024-07-10 06:12:19 -07:00

extend_test.py

remove jax.extend.backend.default_backend in favor of jax.backend

2024-08-02 07:07:29 -07:00

fft_test.py

Follow-up to #22736

2024-07-30 05:39:19 +02:00

for_loop_test.py

Replace jax.xla_computation with the AOT API and add a way to unaccelerate the deprecation in jax tests.

2024-06-18 15:47:24 -07:00

fused_attention_stablehlo_test.py

add sm86/sm89

2024-07-24 21:12:37 +00:00

generated_fun_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

gpu_memory_flags_test.py

Add CUDA custom call example as a JAX test

2024-06-17 15:21:49 -04:00

heap_profiler_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

host_callback_test.py

Remove linear parameter from lax.cond_p.

2024-07-01 10:25:42 -04:00

host_callback_to_tf_test.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

image_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

infeed_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

jax_jit_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jax_to_ir_test.py

Ignore the Deprecation warning produced about native_serialization=False.

2024-07-29 09:11:54 -07:00

jaxpr_effects_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

jaxpr_util_test.py

add pjit forwarding rule

2024-05-25 17:46:01 +00:00

jet_test.py

Updated jnp.ceil/floor/trunc to preserve int dtypes

2024-06-25 20:26:53 +02:00

key_reuse_test.py

[key reuse] handle reuse of closed-over constants

2024-04-11 15:39:45 -07:00

lax_autodiff_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

lax_control_flow_test.py

deflake cond memory leak regression test

2024-07-25 23:12:21 +00:00

lax_metal_test.py

adopt numpy rank promotion config in metal test

2024-07-18 10:57:43 -07:00

lax_numpy_einsum_test.py

jnp.einsum_path: improve docs & annotations

2024-05-10 08:39:32 -07:00

lax_numpy_indexing_test.py

Improve error when indexing with too many indices

2024-07-31 13:57:48 -07:00

lax_numpy_operators_test.py

Added op name and dtype into testcase name for tests in lax_numpy_operators_test.py

2024-05-27 10:31:59 +00:00

lax_numpy_reducers_test.py

Fix compatibility with nightly numpy

2024-07-03 11:07:58 -04:00

lax_numpy_test.py

Add concretization error check in core.min_dim and core.max_dim

2024-08-01 07:27:35 +02:00

lax_numpy_ufuncs_test.py

jnp.vectorize: respect numpy_rank_promotion config

2024-07-08 09:03:03 -07:00

lax_numpy_vectorize_test.py

jnp.vectorize: respect numpy_rank_promotion config

2024-07-08 09:03:03 -07:00

lax_scipy_sparse_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_scipy_special_functions_test.py

Copy nn.{softmax,log_softmax} to scipy.special

2024-06-22 09:32:14 -04:00

lax_scipy_spectral_dac_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

lax_scipy_test.py

Fix or disable some tests that fail when using a Eigen BLAS with AVX vectorization.

2024-07-31 10:06:45 -07:00

lax_test.py

Fix arcsinh accuracy test

2024-07-31 13:23:12 +03:00

lax_vmap_op_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

lax_vmap_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

layout_test.py

Skip test_concrete_layout_in_shardings on GPU backend.

2024-07-30 11:10:29 -07:00

lazy_loader_test.py

lazy_loader_module: Move to new internal_test_util directory.

2023-03-10 10:29:33 -08:00

linalg_test.py

Fix or disable some tests that fail when using a Eigen BLAS with AVX vectorization.

2024-07-31 10:06:45 -07:00

lobpcg_test.py

Removed mentions of iree from the test suite

2024-05-24 10:31:57 +01:00

logging_test.py

jax.config: validate on set()

2024-06-25 09:02:32 -07:00

lru_cache_test.py

Storing the last access time of a cache entry in a separate file

2024-07-03 13:54:36 +04:00

memories_test.py

Fix the AOT check for sharding consistency which skipped checking the devices of the sharding.

2024-08-02 08:15:32 -07:00

mesh_utils_test.py

mesh_utils: allow meshes that do not include device at (0, 0, 0).

2024-07-31 01:11:31 -07:00

metadata_test.py

Disable test_source_file_prefix_removal test because there is cross-contamination of metadata information from different call sites because of cached jitted functions

2024-07-12 12:07:24 -07:00

mock_gpu_test.py

Rename mock_num_processes to mock_num_gpu_processes since this flag is only for GPUs. The naming change was a regression introduced in https://github.com/google/jax/pull/22619

2024-07-31 10:41:03 -07:00

monitoring_test.py

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -07:00

mosaic_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

multi_device_test.py

[ROCM] test_computation_follows_data fix for rocm and cuda

2024-07-15 16:28:22 -05:00

multibackend_test.py

fix some additional warnings related to #21834

2024-06-13 16:06:14 -07:00

multiprocess_gpu_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

mutable_array_test.py

Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file.

2024-05-02 14:36:16 -04:00

name_stack_test.py

Replace xla_computation in name_stack_test with jit(f).lower(...).compiler_ir()

2024-06-12 18:57:20 -07:00

nn_test.py

Support GQA and MQA

2024-07-29 17:17:22 +00:00

ode_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

optimizers_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

package_structure_test.py

Export KeyPath and related types to jax.tree_util

2024-07-31 06:41:33 -07:00

pgle_test.py

[JAX] Fix FDO profile deserialization.

2024-06-20 04:31:34 -07:00

pickle_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

pjit_test.py

Thread the mesh context manager to the place where we recover out_shardings back from GSPMDShardings. Before if you had a program like this:

2024-08-02 11:04:48 -07:00

pmap_test.py

remove pdot and xeinsum (since xmap is gone)

2024-07-25 21:19:17 +00:00

polynomial_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

pretty_printer_test.py

Compute source maps when pretty-printing jaxprs.

2024-05-06 15:45:25 -04:00

profiler_test.py

Support pathlib.Path input to start_trace

2024-07-19 23:41:02 +03:00

python_callback_test.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

pytorch_interoperability_test.py

Enable some more C API tests.

2024-04-22 09:38:59 -07:00

qdwh_test.py

Bump up tolerance for tests using _testQdwh.

2024-06-14 23:18:20 -07:00

random_lax_test.py

Finalize deprecation of batched keys to PRNG functions

2024-05-22 09:40:32 -07:00

random_test.py

make core.as_named_shape treat int like tuple[int]

2024-07-20 17:14:38 +00:00

scipy_fft_test.py

jax.scipy.fft.dct: implement & test norm='backward'

2024-07-22 11:18:35 -07:00

scipy_interpolate_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

scipy_ndimage_test.py

DOC: Improve remaining jax.scipy docstrings

2024-05-02 10:29:00 -07:00

scipy_optimize_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

scipy_signal_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

scipy_spatial_test.py

Add canonical arg to Rotation.as_quat() and switch .inv() to use the quaternion conjugate.

2024-07-24 15:29:19 -07:00

scipy_stats_test.py

Fix betabinom.logpmf and binom.logpmf for JAX to emulate SciPy's behavior when k=n=0

2024-07-31 07:58:43 +05:30

shape_poly_test.py

Add concretization error check in core.min_dim and core.max_dim

2024-08-01 07:27:35 +02:00

shard_alike_test.py

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -07:00

shard_map_test.py

If the product of manual axes is of size 1, then skip emitting

2024-07-31 13:12:06 -07:00

source_info_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

sourcemap_test.py

Merge pull request #21096 from gspschmid:gschmid/sourcemaps

2024-05-08 05:44:08 -07:00

sparse_bcoo_bcsr_test.py

[ROCM] Fix version checks after rocm pjrt integration

2024-07-12 08:07:52 -05:00

sparse_nm_test.py

Disable tests that trigger warning if x64 mode isn't enabled.

2024-05-10 19:58:22 +00:00

sparse_test.py

Re-pack gradients of jax.experimental.sparse.grad() to match original pytrees & test cases

2024-07-29 13:04:05 -04:00

sparsify_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

stack_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

state_test.py

Add some hypothesis testing utilities and developer documentation.

2024-07-15 17:05:32 +02:00

stax_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

svd_test.py

Bump some SVD test tolerances.

2024-07-15 12:54:00 -07:00

transfer_guard_test.py

fix some additional warnings related to #21834

2024-06-13 16:06:14 -07:00

tree_util_test.py

tree_all: add support for is_leaf

2024-06-10 09:46:15 -07:00

typing_test.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

util_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

version_test.py

Add more extensive tests for version strings

2023-09-08 11:33:49 -07:00

x64_context_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

xla_bridge_test.py

Bump minimum jaxlib version to v0.4.30.

2024-06-18 12:35:08 -04:00

xla_interpreter_test.py

Move some utilities out of dispatch.py next to their users, add more types.

2023-08-08 10:52:11 -07:00