Skye Wanderman-Milne a80cbc5626 [JAX] Implement the stream argument to jax.Array.__dlpack__ for CUDA GPU ...

Also implements jax.Array.__dlpack_device__. See
https://dmlc.github.io/dlpack/latest/python_spec.html

This requires plumbing the raw CUDA stream pointer through PJRT and
StreamExecutor (since the GPU PJRT implementation is still based on
SE). This is done via the new PJRT method
ExternalReference::WaitUntilBufferReadyOnStream.

I haven't plumbed this through the PJRT C API yet, because I'm still
debating whether this should be part of the main API or a GPU-specific
extension (plus either way it should probably be its own change).

PiperOrigin-RevId: 558245360

2023-08-18 14:20:38 -07:00

..

filecheck

Fix symbol collision when merging MLIR modules.

2023-06-20 13:53:30 -07:00

notebooks

Delete remote TPU support.

2023-03-24 12:33:33 -07:00

pallas

Skip PallasTest when triton is not installed.

2023-08-16 09:12:51 -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

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

ann_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

aot_test.py

Add internal jaxlib function for fetching the topology from

2023-08-01 14:54:08 -07:00

api_test.py

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

api_util_test.py

jax.jit now works correctly if both donate_argnums and donate_argnames are specified.

2023-07-14 14:28:16 -07:00

array_interoperability_test.py

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

array_test.py

Reverts 7012a05497faf4d33c967bee3cebc83588234e63

2023-08-11 10:30:13 -07:00

batching_test.py

Skip tridiagonal solve batching test on older jaxlibs.

2023-08-11 09:12:03 -07:00

BUILD

Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.

2023-08-15 06:39:46 -07:00

cache_key_test.py

Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.

2023-08-15 06:39:46 -07:00

checkify_test.py

Better pprint rule for check_p primitive.

2023-06-12 10:58:40 -07:00

clear_backends_test.py

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

compilation_cache_test.py

Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.

2023-08-15 06:39:46 -07:00

core_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

custom_linear_solve_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

custom_object_test.py

Move functions out of xla.py closer to their users.

2023-08-08 14:40:42 -07:00

custom_root_test.py

Relax test tolerances.

2023-06-23 09:22:11 -07:00

debug_nans_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

debugger_test.py

Apply pyupgrade --py39-plus.

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

debugging_primitives_test.py

Bump minimum jaxlib version to 0.4.11. xla_extension_version is 158 and mlir_api_version is 49. It will subsume https://github.com/google/jax/pull/16161#issuecomment-1564977332

2023-06-01 09:42:55 -07:00

deprecation_test.py

deprecation_module: Move to new internal_test_util directory.

2023-02-17 10:55:04 -08:00

dtypes_test.py

Add support for float8_e4m3fnuz and float8_e5m2fnuz.

2023-08-07 11:48:53 +01:00

dynamic_api_test.py

Merge pull request #16719 from axch:ragged-test-cases

2023-07-14 15:03:13 -07:00

errors_test.py

Tracebacks no longer have JAX-internal frames prepended by default

2023-08-03 11:38:38 -07:00

experimental_rnn_test.py

Disable the RunnTest.test_lstm1 test since it is fixed for cudnn >= 8.8

2023-05-31 06:21:01 -07:00

fft_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

for_loop_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

generated_fun_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -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

heap_profiler_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

host_callback_test.py

[host_callback] Remove old backwards compatibility flag jax_host_callback_ad_transforms.

2023-08-16 10:01:49 -07:00

host_callback_to_tf_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

image_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

infeed_test.py

Remove 'pjrt_c_api_unimplemented' pytest mark.

2023-03-24 23:14:54 +00:00

jax_jit_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

jax_to_ir_test.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

jaxpr_effects_test.py

Fix scan DCE rule to update effects

2023-05-16 23:17:02 +00:00

jaxpr_util_test.py

Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target.

2023-08-08 10:09:09 -07:00

jet_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_autodiff_test.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

lax_control_flow_test.py

Apply pyupgrade --py39-plus.

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

lax_numpy_einsum_test.py

Reverts 75c3457264f9cc117ff09551ce3174d72689fa3d

2023-08-16 16:06:28 -07:00

lax_numpy_indexing_test.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

lax_numpy_operators_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_numpy_reducers_test.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

lax_numpy_test.py

Re-enable mixed-precision testDot on TPU

2023-08-16 08:23:58 -07:00

lax_numpy_ufuncs_test.py

jnp.ufunc: minor cleanups & test fixes

2023-08-14 15:19:46 -07:00

lax_numpy_vectorize_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_scipy_sparse_test.py

Relax test tolerances.

2023-06-23 09:22:11 -07:00

lax_scipy_special_functions_test.py

Create lax.zeta with native HLO lowering

2023-08-16 13:43:41 -07:00

lax_scipy_spectral_dac_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_scipy_test.py

Create lax.zeta with native HLO lowering

2023-08-16 13:43:41 -07:00

lax_test.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

lax_vmap_op_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

lax_vmap_test.py

Fix scatter batching rule for scatter_apply

2023-07-10 16:42:45 -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

Second attempt at fixing funm tolerance for LLVM change.

2023-08-14 13:01:19 -07:00

lobpcg_test.py

Remove matplotlib from the test requirements.

2023-06-14 09:02:49 -04:00

logging_test.py

Add jax_debug_log_modules config option.

2023-07-28 18:11:12 +00:00

mesh_utils_test.py

Apply pyupgrade --py39-plus.

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

metadata_test.py

Change metadata_test to tolerate paths with backslashes.

2023-06-12 21:57:41 -04:00

monitoring_test.py

Add functions to unregister event duration listeners.

2023-08-09 10:43:02 -07:00

mosaic_test.py

Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic

2023-07-20 18:28:51 -07:00

multi_device_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

multibackend_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

multiprocess_gpu_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

name_stack_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

nn_test.py

jax.nn.softmax: fix fill value when where is specified

2023-06-01 10:18:05 -07:00

ode_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

optimizers_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

package_structure_test.py

Apply pyupgrade --py39-plus.

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

pgle_test.py

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

pickle_test.py

Make inspect_sharding_lowering_rule work only with HloSharding in it's callback. Also remove get_replicated_op_sharding since it is not needed anymore.

2023-06-07 14:39:52 -07:00

pjit_test.py

Force inline all the calls in the module created by a custom call.

2023-08-18 11:13:45 -07:00

pmap_test.py

Move functions out of xla.py closer to their users.

2023-08-08 14:40:42 -07:00

polynomial_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

profiler_test.py

Bump ProfilerTest.test_remote_profiler to run computations for 10 seconds.

2023-07-27 23:54:42 +00:00

python_callback_test.py

Apply pyupgrade --py39-plus.

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

pytorch_interoperability_test.py

[JAX] Implement the stream argument to jax.Array.__dlpack__ for CUDA GPU

2023-08-18 14:20:38 -07:00

qdwh_test.py

Fix test failures in JAX under NumPy 1.25.0rc1.

2023-06-09 14:10:35 -07:00

random_test.py

add random.triangular and random.lognormal

2023-08-05 12:27:53 +08:00

scipy_fft_test.py

Disable testiDctn just like testiDct

2023-07-31 13:58:46 -07:00

scipy_interpolate_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_ndimage_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_optimize_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_signal_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

scipy_spatial_test.py

Fix problematic rotation tests

2023-06-20 05:27:20 -05:00

scipy_stats_test.py

Add jax.scipy.stats.binom

2023-06-27 03:41:38 -07:00

shard_map_test.py

Apply pyupgrade --py39-plus.

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

source_info_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

sparse_test.py

Merge pull request #16975 from hawkinsp:win

2023-08-07 08:10:21 -07:00

sparsify_test.py

Previously, using sparse.todense on a BCSR matrix with sparse.sparsify would raise NotImplementedError: sparse rule for todense is not implemented. By adding the sparse rule, it will resolve this issue.

2023-07-26 13:01:02 -07:00

stack_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

state_test.py

Apply pyupgrade --py39-plus.

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

stax_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

svd_test.py

[ROCm]: Re-enable SvdOnTinyElement test

2023-06-23 17:50:00 +00:00

transfer_guard_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

tree_util_test.py

Apply pyupgrade --py39-plus.

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

typing_test.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

util_test.py

Remove obsolete jaxlib version checks

2023-07-12 11:53:55 -07:00

version_test.py

Check that jaxlib's nightly date is not greater than jax's nightly date in check_jaxlib_version

2023-08-18 10:54:52 -07:00

x64_context_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

xla_bridge_test.py

[PJRT C API] Let framework explicitly check whether a plugin is initialized and initialize the plugin.

2023-08-15 15:24:14 -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

[XLA:Client] Make HloSharding::iota_tile actually produce V2 shardings.

2023-08-07 16:46:53 -07:00