Sergei Lebedev 9584ee3bb9 [pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test ...

Turns out we can mix parallel grid with `plgpu.emit_pipeline` without doing
indexing at all!

PiperOrigin-RevId: 698442820

2024-11-20 10:42:02 -08:00

..

filecheck

Register acos primitive to lower to CHLO acos.

2024-10-25 13:20:36 -07:00

mosaic

[mosaic_gpu] Fixed FragmentedArray comparisons with literals

2024-11-20 04:31:28 -08:00

notebooks

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

pallas

[pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test

2024-11-20 10:42:02 -08:00

testdata

…

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

Make a direct linearize trace.

2024-11-20 10:03:00 -08:00

api_util_test.py

Consolidate handling of input argument resolution in custom_* APIs.

2024-08-13 00:30:23 -07:00

array_api_skips.txt

Fix array API tests.

2024-11-01 09:48:45 -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

Set PYTHONWARNINGS=error in bazel tests.

2024-09-24 12:30:11 -07:00

array_test.py

Make Tracers non-hashable

2024-11-05 09:08:33 -08:00

attrs_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

batching_test.py

Remove UnshapedArray values from JAX (it remains as an abstract class).

2024-10-30 18:53:51 -07:00

blocked_sampler_test.py

[Pallas] Implement block-invariant sampling.

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

BUILD

[SDY] fix JAX layouts tests for Shardy.

2024-11-18 12:14:32 -08:00

cache_key_test.py

Squashed commit of the following:

2024-09-10 17:30:08 -07:00

checkify_test.py

add checkify rule for remat

2024-10-01 02:01:18 +00: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

colocated_python_test.py

Fix ColocatedPythonTest. The test has been failing only on pytest nighlies because there are more GPU devices than CPU devices available, but the tests was making assumption that number of cpu devices is always bigger.

2024-11-03 08:41:16 -08:00

compilation_cache_test.py

Add test utility for accessing jaxlib version tuple.

2024-11-19 12:00:32 -08: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

Reverts f281c6f46475270a57a02416469226315377592c

2024-11-05 07:17:14 -08:00

cudnn_fusion_test.py

Re-enable cudnn_fusion_test on A100.

2024-11-01 15:48:07 +00:00

custom_linear_solve_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

custom_root_test.py

Avoid 'from jax import config' imports

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

debug_nans_test.py

Add a temporary fix for spurious debug_nans errors when round-tripping jaxprs.

2024-10-30 11:34:08 -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

Make debug.print work with static args. Fixes: https://github.com/google/jax/issues/23600

2024-09-18 08:41:29 -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

Add float8_e4m3 and float8_e3m4 types support

2024-11-08 18:58:31 +00:00

dynamic_api_test.py

simplify conversion logic involving extended dtypes

2024-09-25 00:10:01 +00:00

errors_test.py

Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.

2024-09-26 08:39:30 -07:00

experimental_rnn_test.py

…

export_back_compat_test.py

Activate the FFI implementation of SVD on GPU.

2024-10-17 17:57:06 -07:00

export_harnesses_multi_platform_test.py

[export] Cleanup in the export module.

2024-11-01 22:56:44 -07:00

export_test.py

[export] Cleanup in the export module.

2024-11-01 22:56:44 -07:00

extend_test.py

Merge pull request #24593 from froystig:random-dtypes

2024-11-19 23:04:06 -08:00

fft_test.py

Remove checks for jaxlib v0.4.33 in tests

2024-10-11 15:39:24 -04:00

for_loop_test.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00

fused_attention_stablehlo_test.py

fix bwd batcher for unsupported dbias

2024-09-23 17:43:25 +00:00

garbage_collection_guard_test.py

Merge pull request #24828 from jax-ml:caching-fix

2024-11-11 10:23:51 -08: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

image_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

infeed_test.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00

jax_jit_test.py

…

jax_to_ir_test.py

[jax2tf] Remove non-native serialization test from jax_to_ir_test

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

Adding support for copy_p primitive to jet.

2024-11-07 17:06:02 -08: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

Reverts 2075b091c4e83f0bdbd0d47812a72114fb8b937a

2024-11-19 14:47:24 -08:00

lax_control_flow_test.py

Change the lowering rule for jax.lax.scan to avoid emitting a while loop

2024-10-30 06:20:39 -07:00

lax_metal_test.py

adopt jax.extend.backend

2024-10-07 12:27:35 -07:00

lax_numpy_einsum_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

lax_numpy_indexing_test.py

Added scatter_sub_p

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

lax_numpy_operators_test.py

Implement jax.numpy.spacing

2024-10-03 10:40:39 -07:00

lax_numpy_reducers_test.py

[tpu] Disable a cumulative reduction test on TPU v6e that currently hits an unimplemented case in XLA.

2024-11-04 08:39:46 -08:00

lax_numpy_test.py

Add missing version guard in GPU tests for jnp.poly.

2024-11-19 09:52:45 -08:00

lax_numpy_ufuncs_test.py

Make logaddexp and logaddexp2 into ufuncs

2024-11-18 09:27:36 -08:00

lax_numpy_vectorize_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

lax_scipy_sparse_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

lax_scipy_special_functions_test.py

Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.

2024-11-18 20:33:27 -05: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

Tighten test tolerances after the underlying issue causing nondeterministic results for _nrm2 in Eigen BLAS was fixed in https://gitlab.com/libeigen/eigen/-/merge_requests/1667 -> cl/663346025

2024-09-20 10:03:46 -07:00

lax_test.py

Reverts 2075b091c4e83f0bdbd0d47812a72114fb8b937a

2024-11-19 14:47:24 -08:00

lax_vmap_op_test.py

Avoid 'from jax import config' imports

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

lax_vmap_test.py

Reverts 2075b091c4e83f0bdbd0d47812a72114fb8b937a

2024-11-19 14:47:24 -08:00

layout_test.py

[XLA] Make our LLVM usage more googley

2024-11-12 22:00:38 -08:00

lazy_loader_test.py

…

linalg_test.py

Add test utility for accessing jaxlib version tuple.

2024-11-19 12:00:32 -08:00

lobpcg_test.py

Removed mentions of iree from the test suite

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

logging_test.py

logging fixes

2024-11-05 13:28:17 -08: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

magma_linalg_test.py

Add test utility for accessing jaxlib version tuple.

2024-11-19 12:00:32 -08:00

memories_test.py

Add test_compute_on_host_shared_sharding in memories_test

2024-11-19 21:33:27 -08:00

mesh_utils_test.py

Removed mesh_utils._bounds_from_last_device which was only used in tests

2024-10-30 02:43:56 -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

#sdy add shardy CPU config for all JAX tests, disabling any known failing test cases.

2024-10-30 11:40:20 -07:00

mock_gpu_topology_test.py

Add an option to specify mock GPU topology

2024-11-12 08:36:27 -08:00

monitoring_test.py

…

mosaic_test.py

Avoid 'from jax import config' imports

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

multi_device_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

multibackend_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

multiprocess_gpu_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

mutable_array_test.py

Fix mutable array effects not being tracked properly

2024-09-30 18:55:15 -07: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

Adjusted atol/rtol for jax sdpa tests

2024-10-21 17:00:04 +00:00

ode_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

optimizers_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

package_structure_test.py

Set __module__ attribute for objects in jax.numpy

2024-11-15 06:03:54 -08:00

pgle_test.py

[AutoPGLE] Use compile options to override debug options instead of XLA_FLAGS.

2024-11-19 01:47:54 -08:00

pickle_test.py

Avoid 'from jax import config' imports

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

pjit_test.py

Fix a bug where constant deduplication used an inappropriate inequality.

2024-11-19 18:59:49 -08:00

pmap_test.py

Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified.

2024-11-07 15:51:54 -08: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

Rename vmap methods for callbacks.

2024-10-21 15:03:04 -04:00

pytorch_interoperability_test.py

Set PYTHONWARNINGS=error in bazel tests.

2024-09-24 12:30:11 -07:00

qdwh_test.py

Bump up tolerance for tests using _testQdwh.

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

random_lax_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

random_test.py

Merge pull request #24593 from froystig:random-dtypes

2024-11-19 23:04:06 -08:00

scipy_fft_test.py

Fix dtype bug in jax.scipy.fft.idct

2024-09-25 12:55:43 -04:00

scipy_interpolate_test.py

Avoid 'from jax import config' imports

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

scipy_ndimage_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

scipy_optimize_test.py

Fix some warnings causing CI failures on ARM.

2024-09-24 17:25:26 -07:00

scipy_signal_test.py

Avoid 'from jax import config' imports

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

scipy_spatial_test.py

Don't perform size 0 slices into scipy rotations.

2024-11-08 09:56:12 -05:00

scipy_stats_test.py

Remove warning filters from pyproject.toml, add local warning

2024-09-24 01:38:24 +00:00

shape_poly_test.py

[shape_poly] Fix the handling of jvp(lax.sort)

2024-11-12 03:36:05 -08:00

shard_alike_test.py

Remove unused import

2024-11-08 13:36:28 -08:00

shard_map_test.py

Don't psum over auto mesh dims in _unmentioned2.

2024-11-20 10:36:03 -08: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

Set PYTHONWARNINGS=error in bazel tests.

2024-09-24 12:30:11 -07:00

sparse_nm_test.py

[XLA:GPU] Skip small tile sizes for sparse gemms on Ampere as well. Enable the JAX test again that has been failing.

2024-11-11 08:58:15 -08:00

sparse_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

sparsify_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

stack_test.py

Avoid 'from jax import config' imports

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

state_test.py

Reverts f281c6f46475270a57a02416469226315377592c

2024-11-05 07:17:14 -08:00

stax_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -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

register_dataclass: allow marking static fields via field(static=True)

2024-11-06 11:18:11 -08:00

typing_test.py

Upgrade remaining sources to Python 3.9

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

util_test.py

Stop using generators for linear_util transformations.

2024-11-13 13:47:07 -08:00

version_test.py

Adds an env that can let users provide a custom version suffix for jax dev build.

2024-11-16 14:39:54 -08:00

x64_context_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

xla_bridge_test.py

Remove classic HLO lowering rule support from JAX.

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

xla_interpreter_test.py

…

xla_metadata_test.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00