jax authors cb0c49850c Merge pull request #21081 from hawkinsp:sourcemap ...

PiperOrigin-RevId: 631236806

2024-05-06 17:33:12 -07:00

..

filecheck

tests: access tree utilities via jax.tree.*

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

mosaic

[Mosaic GPU] Add the flash attention example

2024-04-29 09:31:30 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

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

pallas

[export] Add backwards compatibility test for Pallas call on GPUs.

2024-05-02 05:38: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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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 donation with kwargs. The problem is that pytrees sort dictionaries by default. So if we create the donation vector with original kwargs order, it won't match the aval order (which is created by sorting kwargs i.e. dict) and we end up donating the wrong input.

2024-04-23 14:50:04 -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

Add get_device_ordinal to cuda plugin so that CUDA dependency can be removed from py_array (jaxlib).

2024-04-30 12:50:50 -07:00

array_test.py

If callback returns a fully replicated global array, return it as is.

2024-04-14 14:35:57 -07:00

attrs_test.py

[attrs] fix tracer lifetime bug, fixes #20082

2024-03-05 12:08:44 -08:00

batching_test.py

Avoid 'from jax import config' imports

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

BUILD

Compute source maps when pretty-printing jaxprs.

2024-05-06 15:45:25 -04: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

clear_backends_test.py

Avoid 'from jax import config' imports

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

compilation_cache_test.py

Fix warnings in CI from compilation_cache_test.

2024-04-30 13:40:35 +00:00

core_test.py

Allow replacing jaxpr debug_info with None.

2024-04-30 08:31:39 -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

Avoid 'from jax import config' imports

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

debugger_test.py

jax.debug.callback now passes arguments as jax.Arrays

2024-04-19 13:57:18 -07:00

debugging_primitives_test.py

jax.debug.callback now passes arguments as jax.Arrays

2024-04-19 13:57:18 -07:00

deprecation_test.py

Fix tests that mistakenly used assertRaises(..., msg=...) to match a message.

2023-10-06 08:07:49 -04:00

dtypes_test.py

Added int4 and uint4 to dtype-specific tests

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

dynamic_api_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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] Add backwards compatibility test for Pallas call on GPUs.

2024-05-02 05:38:33 -07:00

export_harnesses_multi_platform_test.py

[ROCm]: Add rocm as platform in export_harnessess_multi_platform_test

2024-03-20 16:24:06 +00:00

export_test.py

Fix mlir.merge_mlir_modules to properly remember the inlined symbols

2024-04-29 19:56:41 +03:00

extend_test.py

Avoid 'from jax import config' imports

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

fft_test.py

Updated

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

for_loop_test.py

Avoid 'from jax import config' imports

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

fused_attention_stablehlo_test.py

Support BNTH input formats

2024-04-03 20:48:37 +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

Exclude test_gpu_memory_allocation from pytest execution.

2023-12-20 14:53:33 -08:00

heap_profiler_test.py

Avoid 'from jax import config' imports

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

host_callback_test.py

Avoid 'from jax import config' imports

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

Made has_side_effect= parameter of mlir.emit_python_callback keyword-only

2024-05-02 13:44:54 -07:00

jaxpr_util_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

Add specialize on jax.jit so that we can delete the duplicate code in jax.make_jaxpr.

2024-04-27 18:58:16 -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

Avoid 'from jax import config' imports

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

lax_numpy_indexing_test.py

jax.numpy: implement scalar boolean indexing

2024-02-09 11:00:00 -08:00

lax_numpy_operators_test.py

Fixed hypot bug on nan/inf pairings, began deprecation of non-real values

2024-04-15 17:56:16 +00:00

lax_numpy_reducers_test.py

Refactored common upcast for integral-type accumulators

2024-05-06 15:13:10 +00:00

lax_numpy_test.py

Refactor jax.numpy docstring tests

2024-05-03 11:04:43 -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

Add jax.scipy.special.gammasgn

2024-04-18 16:14:55 -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

Avoid 'from jax import config' imports

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

lax_test.py

Workaround mpmath 1.3 bugs in tan and tanh evaluation at infinities

2024-04-10 18:26:07 +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

Add layout support to make_array_from_callback.

2024-04-15 12:38:34 -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

jnp.linalg tensorinv & tensorsolve: improve implementation & docs

2024-05-06 11:08:36 -07:00

lobpcg_test.py

Avoid 'from jax import config' imports

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

logging_test.py

Avoid 'from jax import config' imports

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

memories_test.py

[XLA:CPU] Enable constant host offloading

2024-04-30 14:51:11 -07:00

mesh_utils_test.py

Adding expectations to e2e test with expectations for mesh creation allowing splitting physical axes.

2024-03-21 18:58:06 -07:00

metadata_test.py

Avoid 'from jax import config' imports

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

mock_gpu_test.py

Avoid 'from jax import config' imports

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

Avoid 'from jax import config' imports

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

multibackend_test.py

Avoid 'from jax import config' imports

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

multiprocess_gpu_test.py

Avoid 'from jax import config' imports

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

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

pgle_test.py

Avoid 'from jax import config' imports

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

pickle_test.py

Avoid 'from jax import config' imports

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

pjit_test.py

Add specialize on jax.jit so that we can delete the duplicate code in jax.make_jaxpr.

2024-04-27 18:58:16 -07:00

pmap_test.py

Guard test that requires two devices with device_count() check.

2024-04-01 12:32:54 -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

relax a side-effects test that was erroneously checking against a canonical order

2024-04-15 17:13:35 -07:00

pytorch_interoperability_test.py

Enable some more C API tests.

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

qdwh_test.py

Refactor QDWH to be more efficient when run batched under vmap.

2024-04-25 11:48:21 -07:00

random_lax_test.py

random_lax_test: fix kstest for newer NumPy

2024-05-06 09:20:07 -07:00

random_test.py

reintroduce the Threefry GPU kernel lowering, under a flag

2024-05-01 10:33:31 -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

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

shard_alike_test.py

Avoid 'from jax import config' imports

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

shard_map_test.py

[shard_map] better fix for spmd_axis_name issues with shmap residuals

2024-05-04 01:31:15 +00:00

source_info_test.py

Avoid 'from jax import config' imports

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

Use correct kWidth in sparse dots with int8 input (on Ampere)

2024-04-26 04:53:25 -07: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

Avoid 'from jax import config' imports

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

tree_util_test.py

Add builtin cc dataclass pytree node for performance.

2024-04-23 14:14:49 -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

Change xla_bridge_test to expect a bytes FDO profile instead of a string in future jaxlib versions.

2024-03-06 13:50:36 -08: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

Avoid 'from jax import config' imports

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