Yash Katariya 6007698f4e Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated. ...

The semantics are as follow:

* if the mesh context manager is not provided, None will be treated as UNSPECIFIED for both in_shardings and out_shardings

* If the mesh context manager is provided, None will be treated as fully replicated as per the old semantics.

This will make sure that we don't break existing code depending on None meaning replicated but also start making the transition to None meaning UNSPECIFIED for jit and pjit.

PiperOrigin-RevId: 540705660

2023-06-15 15:22:22 -07:00

..

filecheck

Migrate regularized_incomplete_beta_p off xla_fallback

2023-03-24 14:53:20 -07:00

notebooks

Delete remote TPU support.

2023-03-24 12:33:33 -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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

ann_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

aot_test.py

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00

api_test.py

Added argument jax.linearize(..., has_aux=...)

2023-06-14 22:34:13 -07:00

api_util_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

array_interoperability_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

array_test.py

Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated.

2023-06-15 15:22:22 -07:00

batching_test.py

Fix some test failures under H100.

2023-06-09 09:23:36 -07:00

BUILD

Added scipy.spatial.transform Rotation and Slerp classes

2023-06-08 07:51:32 -05:00

checkify_test.py

Better pprint rule for check_p primitive.

2023-06-12 10:58:40 -07:00

clear_backends_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

compilation_cache_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

core_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

custom_linear_solve_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

custom_object_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

custom_root_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

debug_nans_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

debugger_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

debugging_primitives_test.py

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00

deprecation_test.py

deprecation_module: Move to new internal_test_util directory.

2023-02-17 10:55:04 -08:00

dtypes_test.py

Clean up old jaxlib version guards.

2023-06-06 01:17:30 -07:00

dynamic_api_test.py

Test and fix jnp.broadcast_to.

2023-06-08 16:17:43 -04:00

errors_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

experimental_rnn_test.py

Disable the RunnTest.test_lstm1 test since it is fixed for cudnn >= 8.8

2023-05-31 06:21:01 -07:00

fft_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

for_loop_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

generated_fun_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

gfile_cache_test.py

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

heap_profiler_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

host_callback_test.py

resolve opaque dtypes in MLIR callback lowering and in XLA shape translation

2023-05-01 08:21:54 -07:00

host_callback_to_tf_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

image_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

infeed_test.py

Remove 'pjrt_c_api_unimplemented' pytest mark.

2023-03-24 23:14:54 +00:00

jax_jit_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

jax_to_ir_test.py

Delete jax_experimental_name_stack flag

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

jaxpr_effects_test.py

Fix scan DCE rule to update effects

2023-05-16 23:17:02 +00:00

jaxpr_util_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

jet_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_autodiff_test.py

allow lax.dot_general to accept different input dtypes

2023-05-22 10:33:42 -07:00

lax_control_flow_test.py

remove unused variable in lax.scan test

2023-06-10 14:05:12 +01:00

lax_numpy_einsum_test.py

Fix some test failures under H100.

2023-06-09 09:23:36 -07:00

lax_numpy_indexing_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_numpy_operators_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_numpy_reducers_test.py

explicitly convert jnp.var scalar normalizer to float (from int)

2023-05-23 09:44:08 -07:00

lax_numpy_test.py

Fix test failure in LaxBackedNumpyTest.testFrexp4 on Windows.

2023-06-13 20:46:06 -04:00

lax_numpy_vectorize_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_scipy_sparse_test.py

Fix some test failures under H100.

2023-06-09 09:23:36 -07:00

lax_scipy_special_functions_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_scipy_spectral_dac_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_scipy_test.py

Add test of private xlogx gradient

2023-04-25 14:31:00 -07:00

lax_test.py

Simplify sharding types input to physical_hlo_sharding and lower_jaxpr_to_fun.

2023-06-06 18:03:12 -07:00

lax_vmap_op_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_vmap_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -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

[XLA:TPU] Fix a bug in eigh that caused a slight loss of accuracy.

2023-05-04 07:49:04 -07:00

lobpcg_test.py

Remove matplotlib from the test requirements.

2023-06-14 09:02:49 -04:00

mesh_utils_test.py

Enable more mesh shape assignment

2023-03-31 09:36:16 -07:00

metadata_test.py

Change metadata_test to tolerate paths with backslashes.

2023-06-12 21:57:41 -04: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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

multibackend_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

multiprocess_gpu_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

name_stack_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

nn_test.py

jax.nn.softmax: fix fill value when where is specified

2023-06-01 10:18:05 -07:00

ode_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

optimizers_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

package_structure_test.py

Add test framework for module attribute

2023-04-21 13:20:16 -07:00

pickle_test.py

Make inspect_sharding_lowering_rule work only with HloSharding in it's callback. Also remove get_replicated_op_sharding since it is not needed anymore.

2023-06-07 14:39:52 -07:00

pjit_test.py

Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated.

2023-06-15 15:22:22 -07:00

pmap_test.py

Fix some test failures under H100.

2023-06-09 09:23:36 -07:00

polynomial_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

profiler_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

python_callback_test.py

Add maximal sharding for pure_callback not inside of a shard_map

2023-05-11 13:28:37 -07:00

pytorch_interoperability_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

qdwh_test.py

Fix test failures in JAX under NumPy 1.25.0rc1.

2023-06-09 14:10:35 -07:00

random_test.py

custom prng: add shard methods to PRNGKeyArrayImpl

2023-06-01 04:10:12 -07:00

scipy_fft_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_interpolate_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_ndimage_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_optimize_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_signal_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_spatial_test.py

Increase precision in LaxBackedScipySpatialTransformTests.testRotationApply

2023-06-08 16:30:58 -07:00

scipy_stats_test.py

Temporarily disable LaxBackedScipyStatsTests.testTruncnormPdf.

2023-05-15 06:52:46 -07:00

shard_map_test.py

fix key opaque dtype physical HLO sharding handler

2023-06-06 15:49:03 -07:00

source_info_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

sparse_test.py

Make dimensions argument of bcoo_reshape optional.

2023-05-09 10:38:18 -07:00

sparsify_test.py

Change np.prod->math.prod

2023-04-13 11:48:11 -07:00

stack_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

state_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

stax_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

svd_test.py

Fix corner cases in JAX SVD: a) Clamp negative singular values to zero. b) Return all NaN for matrices with non-finite values.

2023-06-13 11:06:49 -07:00

transfer_guard_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

tree_util_test.py

Add jax.tree_util.tree_leaves_with_path(tree).

2023-06-12 04:13:37 -07:00

typing_test.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

util_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07: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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

xla_bridge_test.py

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07: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

Fix some test failures under H100.

2023-06-09 09:23:36 -07:00