Sergei Lebedev 575ba942e0 Removed get_compute_capability from jax.experimental.pallas.gpu ...

Compute capability is available as a `str` attribute on a GPU device since
jaxlib 0.4.26.

2024-05-08 21:10:43 +01:00

..

filecheck

tests: access tree utilities via jax.tree.*

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

mosaic

[jax:mosaic-gpu] FragmentedArray can do tiled load.

2024-05-07 18:13:55 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

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

pallas

Removed get_compute_capability from jax.experimental.pallas.gpu

2024-05-08 21:10:43 +01: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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

Merge pull request #21128 from hawkinsp:loggingtest

2024-05-08 10:07:42 -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

clear_backends_test.py

Avoid 'from jax import config' imports

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

compilation_cache_test.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07: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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07: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

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

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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

Handle ellipsis ... in _attempt_rewriting_take_via_slice.

2024-05-07 10:30:08 -07: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

Add optional size argument to jnp.compress & jnp.extract.

2024-05-07 08:47:34 -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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

Bump minimum jaxlib version to 0.4.27

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

pmap_test.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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

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

Removed get_compute_capability from jax.experimental.pallas.gpu

2024-05-08 21:10:43 +01: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

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -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 0.4.27

2024-05-07 16:07:59 -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

Avoid 'from jax import config' imports

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