jax authors 95e2c17b61 Update test of QDWH to use stricter tolerances and test more shapes and types. ...

Get rid of comparison with scipy.linalg.polar, since its outputs are significantly less accurate than QDWH. Since the polar decomposition is unique, comparing to a less accurate implementation does not add value.

PiperOrigin-RevId: 642423757

2024-06-11 16:04:38 -07:00

..

filecheck

tests: access tree utilities via jax.tree.*

2024-02-26 14:17:18 -08:00

mosaic

[Mosaic GPU] Load LLVM lowering interfaces for all dialects

2024-06-10 05:55:01 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

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

pallas

[Pallas] Add pallas TPU random key impls and lowering rules for basic prng ops (seed/foldin/bits/unwrap/wrap).

2024-06-10 18:08:19 -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

Raise an error if device_put sees an invalid value.

2024-06-10 16:07:44 -07: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_test.py

Refactor array_api namespace, relying more directly on jax.numpy

2024-05-02 18:17:45 +00:00

array_interoperability_test.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07:00

array_test.py

Add util.cache to jax.clear_caches and move pjit, sharding, array, etc uses of functools.lru_cache to util.cache so that those caches will be cleared if jax.clear_caches is called.

2024-06-11 12:46:47 -07:00

attrs_test.py

Make attrs work with pytrees

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

batching_test.py

Avoid 'from jax import config' imports

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

BUILD

Update test of QDWH to use stricter tolerances and test more shapes and types.

2024-06-11 16:04:38 -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

Cholesky rank-1 update kernel for JAX.

2024-05-14 15:21:38 -07:00

clear_backends_test.py

Avoid 'from jax import config' imports

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

compilation_cache_test.py

Allow disabling compilation cache for particular runtime_types.

2024-06-04 13:30:42 -07:00

core_test.py

Reverts 1956ff7d7b73794012fece2d8452e097196587fc

2024-05-08 17:23:13 -07: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

tests: set configs with jtu.with_config rather than manually

2024-06-05 13:34:32 -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

test cleanup: use ExitStack to reduce test boilerplate

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

deprecation_test.py

Refactor & test internal deprecation APIs

2024-05-20 17:16:31 -07:00

dtypes_test.py

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

2024-06-03 14:52:50 -07:00

dynamic_api_test.py

Merge pull request #21414 from mattjj:pjit-forwarding-rule

2024-05-25 11:57:07 -07: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

[export] Move the export implementation to jax._src.export.

2024-06-09 08:59:50 -07:00

export_harnesses_multi_platform_test.py

[export] Create the jax.export module APIs.

2024-06-10 19:31:51 +02:00

export_test.py

[export] Cleaned up types of [in|out]_shardings

2024-06-11 13:46:44 +02:00

extend_test.py

Add helper function for building custom call lowering rules

2024-06-06 11:34:08 -04:00

fft_test.py

Updated

2024-02-05 18:01:48 -05:00

for_loop_test.py

Reverts 49bd4d6f01d6cda00f9b1bdfbda156636baae928

2024-05-13 08:35:40 -07:00

fused_attention_stablehlo_test.py

init

2024-05-24 22:26:20 +00:00

generated_fun_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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

gpu_memory_flags_test.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07:00

heap_profiler_test.py

Avoid 'from jax import config' imports

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

host_callback_test.py

test cleanup: use ExitStack to reduce test boilerplate

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

host_callback_to_tf_test.py

[callback] Add a flag to implement host_callback in terms of io_callback.

2024-04-05 08:51:30 +01: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

Added int4 and uint4 to dtype-specific tests

2024-04-18 15:20:20 +01:00

jaxpr_effects_test.py

test cleanup: use ExitStack to reduce test boilerplate

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

jaxpr_util_test.py

add pjit forwarding rule

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

jet_test.py

jet_test: use float32 matmul precision

2024-05-02 14:48:22 +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

Merge code between make_jaxpr and jit(f).trace.

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

lax_metal_test.py

metal_plugin ci with jaxlib nightly

2024-05-02 18:40:01 -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

Support scalar boolean indices in arr.at[idx].set(vals)

2024-05-20 05:33:36 -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

Merge pull request #21262 from vfdev-5:depr-change-ddof-to-correction-21088

2024-05-24 09:47:27 -07:00

lax_numpy_test.py

Updating jnp.abs/absolute docs

2024-06-05 11:27:24 -07:00

lax_numpy_ufuncs_test.py

Avoid 'from jax import config' imports

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

lax_numpy_vectorize_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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

jax.scipy.special.beta: deprecate x,y in favor of a,b

2024-06-10 09:01:39 -07: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

TST: fix Lpmn test for new scipy

2024-05-08 15:55:20 -07:00

lax_test.py

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

2024-06-03 14:52:50 -07: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

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -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

Added trace alias to jnp.linalg

2024-05-29 22:28:44 +00:00

lobpcg_test.py

Removed mentions of iree from the test suite

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

logging_test.py

logging_test: avoid modifying global config state

2024-05-29 13:38:33 -07:00

lru_cache_test.py

Implement LRU cache eviction for persistent compilation cache

2024-06-11 21:48:35 +04:00

memories_test.py

Adds test_compute_no_inputs_host_replicated in memories_test.py

2024-06-10 15:02:34 -07:00

mesh_utils_test.py

Reverts dfe61285093ff826e1ad23bb36b77a42c01040b4

2024-06-06 12:41:17 -07:00

metadata_test.py

Avoid 'from jax import config' imports

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

mock_gpu_test.py

tests: set configs with jtu.with_config rather than manually

2024-06-05 13:34:32 -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

test cleanup: use ExitStack to reduce test boilerplate

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

multibackend_test.py

Avoid 'from jax import config' imports

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

multiprocess_gpu_test.py

tests: set configs with jtu.with_config rather than manually

2024-06-05 13:34:32 -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

Avoid 'from jax import config' imports

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

nn_test.py

softmax: deprecate initial argument & always set to -inf internally

2024-04-10 10:23:21 -07: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

Better documentation for jax.tree_util

2024-05-20 19:56:47 -07:00

pgle_test.py

Disable PgleTest.testPGLEProfilerGetFDOProfile.

2024-05-31 05:40:04 -07:00

pickle_test.py

Avoid 'from jax import config' imports

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

pjit_test.py

Allow generation of sharding strategies with mixed mesh shapes by default.

2024-06-10 09:38:39 -07:00

pmap_test.py

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -07: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

Avoid 'from jax import config' imports

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

python_callback_test.py

test cleanup: use ExitStack to reduce test boilerplate

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

pytorch_interoperability_test.py

Enable some more C API tests.

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

qdwh_test.py

Update test of QDWH to use stricter tolerances and test more shapes and types.

2024-06-11 16:04:38 -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

[Pallas] Add pallas TPU random key impls and lowering rules for basic prng ops (seed/foldin/bits/unwrap/wrap).

2024-06-10 18:08:19 -07:00

scipy_fft_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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

Avoid 'from jax import config' imports

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

scipy_stats_test.py

Merge pull request #20885 from rajasekharporeddy:test_branch4

2024-04-23 13:29:40 -07:00

shape_poly_test.py

Swap operands of dot if the LHS is fed by a parameter

2024-06-10 18:33:05 -07: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

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07: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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07: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 discharge rules for scan with mutable arrays. Move mutable array tests to separate file.

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

stax_test.py

Avoid 'from jax import config' imports

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

svd_test.py

Fix bug in rank-deficient fix-up code: Do not zero out the corresponding column of u_out if a diagonal entry of r is exactly zero.

2024-04-18 09:20:48 -07:00

transfer_guard_test.py

tests: set configs with jtu.with_config rather than manually

2024-06-05 13:34:32 -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

add xla_bridge test guard on cloud tpu env

2024-06-04 13:45:52 -07: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

xmap_test.py

test cleanup: use ExitStack to reduce test boilerplate

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