rocm_jax

filecheck

Integrate StableHLO at openxla/stablehlo@8817ff1d

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

mosaic

Added another boxDim check to mosaic_gpu_init_tma_desc

2024-08-07 03:16:54 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

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

pallas

Merge pull request #22820 from Rifur13:mha-faster

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

Small tweaks to custom_vmap UI.

2024-08-05 17:51:47 +01: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

Do not run nn_test under asan -- it times out

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

Merge pull request #22649 from ROCm:ci_jax_export_harness

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

improve while_loop carry pytree/type mismatch errors

2024-08-03 21:57:29 +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

Support empty boolean indexing

2024-08-06 09:56:03 -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

Rolling forward #22836

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

lax_test.py

Update complex functions accuracy tests

2024-08-02 23:31:51 +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

Add mesh_shape to the lowering context. This is to allow custom partitioning to not depend on the mesh context manager to return NamedShardings even if the arguments have NamedShardings on them.

2024-08-06 18:35:44 -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

Init commit

2024-08-01 19:39:34 +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

Set jax_enable_memories flag to True by default

2024-08-07 16:25:25 -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