rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

2020-06-01 11:49:35 -07:00

filecheck

Lower bessel_i1e primitive to chlo.bessel_i1e.

2022-08-16 12:39:28 -07:00

notebooks

Update colab CPU test notebook to be more robust

2021-07-13 09:30:56 -07:00

third_party/scipy

[x64] make scipy_optimize_test compatible with strict dtype promotion

2022-06-21 09:29:08 -07:00

ann_test.py

[JAX] Fix batch logic for approx_min/max_k

2022-04-08 13:50:36 -07:00

api_test.py

Add jax.print_environment_info()

2022-09-12 15:39:33 -07:00

api_util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

array_interoperability_test.py

Disabled dlpack bfloat16 test for JAX<->cupy.

2022-08-19 10:08:07 -04:00

array_test.py

jax_array_test: set config once & fix X64 failure

2022-09-13 14:06:38 -07:00

batching_test.py

roll-forward #11952 , take 2

2022-08-22 13:57:31 -07:00

BUILD

Use cases_from_list to subsample enumerated cases in for_loop_test

2022-09-13 12:34:10 -07:00

callback_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

checkify_test.py

Checkify: add jvp rule for assert_p.

2022-09-08 05:26:02 -07:00

clear_backends_test.py

Bump the minimum jaxlib version to 0.3.15.

2022-09-08 16:43:46 -04:00

compilation_cache_test.py

[xla:jitrt] Enable serializing GPU executable for JAX

2022-09-08 11:50:31 -07:00

core_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

custom_linear_solve_test.py

add test for custom_linear_solve + new remat

2022-07-28 21:37:23 -07:00

custom_object_test.py

fix custom_object_test for new Array type

2022-08-18 16:33:05 -07:00

custom_root_test.py

Relax tolerances slightly for MKL.

2022-02-25 22:02:55 +00:00

debug_nans_test.py

Fix Forward. The fix is on the user's end. Original PR: https://github.com/google/jax/pull/12217

2022-09-08 08:49:40 -07:00

debugger_test.py

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -07:00

debugging_primitives_test.py

jax.test_util: add capture_stdout context manager

2022-09-12 15:21:52 -07:00

dtypes_test.py

Add array test counterparts to dtypes_test.

2022-08-18 15:58:39 -07:00

dynamic_api_test.py

[dynamic-shapes] add basic vmap-of-indexing support

2022-09-08 17:52:12 -07:00

errors_test.py

Avoid more direct references to jax._src without imports.

2022-08-24 07:51:28 -07:00

fft_test.py

Add shape checks for lax.fft.

2022-08-13 16:32:52 +00:00

for_loop_test.py

Add support for closing over Refs in nested for loops

2022-09-13 13:32:44 -07:00

generated_fun_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

gfile_cache_test.py

Using etils(gfile) to support gcs buckets and file system for persistent compilation caching

2022-06-10 00:17:13 +00:00

global_device_array_test.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

heap_profiler_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

host_callback_test.py

delete old remat implementation

2022-08-16 23:16:37 -07:00

host_callback_to_tf_test.py

Internal change

2022-08-18 11:39:07 -07:00

image_test.py

Internal change

2022-08-18 11:39:07 -07:00

infeed_test.py

Avoid top-level aliases of jax.tree_util.*

2022-07-07 11:41:02 -07:00

jax_jit_test.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

jax_to_ir_test.py

Internal change

2022-08-18 11:39:07 -07:00

jaxpr_effects_test.py

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -07:00

jaxpr_util_test.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

jet_test.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

lax_autodiff_test.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

lax_control_flow_test.py

Add support for closing over Refs in nested for loops

2022-09-13 13:32:44 -07:00

lax_numpy_einsum_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

lax_numpy_indexing_test.py

Copybara import of the project:

2022-08-25 15:14:33 -07:00

lax_numpy_test.py

Enable testArrayCopy now that its fixed.

2022-09-08 14:40:06 -07:00

lax_numpy_vectorize_test.py

[x64] make lax_numpy_vectorize_test work with strict dtype promotion

2022-06-16 14:06:24 -07:00

lax_scipy_sparse_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

lax_scipy_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

lax_test.py

jax.lax: ensure GatherDimensionNumbers contains tuples for hashability

2022-09-12 12:10:17 -07:00

lax_vmap_test.py

Merge pull request #12122 from hawkinsp:fft

2022-08-26 11:32:07 -07:00

linalg_test.py

Bump test tolerance for eigenvalue test.

2022-08-16 14:44:16 -04:00

lobpcg_test.py

Make LOBPCG test plots compatible with bazel.

2022-08-04 20:05:53 -07:00

mesh_utils_test.py

[mesh_utils] Avoid relying on process-tiled device order

2022-06-06 13:30:56 -07:00

metadata_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

multi_device_test.py

Minimally support device argument on jit in the jax.Array path

2022-09-09 16:56:43 -07:00

multibackend_test.py

Minimally support device argument on jit in the jax.Array path

2022-09-09 16:56:43 -07:00

multiprocess_gpu_test.py

Add pytest marker inside the test only if pytest is present in the env

2022-09-06 11:45:59 -07:00

name_stack_test.py

Only wrap one level of the name-stack with transform names.

2022-06-10 03:50:26 -07:00

nn_test.py

Update nn_test.py

2022-08-11 15:46:34 +01:00

ode_test.py

[x64] make ode_test pass with strict dtype promotion

2022-06-17 15:21:10 -07:00

optimizers_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

pickle_test.py

support key array pickling

2022-08-31 12:03:53 -07:00

pjit_test.py

Skip the infeed test since its failing on GPU with jax at HEAD and latest jaxlib on pypi

2022-09-09 20:18:45 -07:00

pmap_test.py

Make __getitem__ work for PmapSharding just like SDA works. DA is already covered with the current implementation.

2022-09-09 14:25:22 -07:00

polynomial_test.py

Remove _ prefix from functions in jax._src.dtypes.

2022-08-12 12:51:09 +00:00

profiler_test.py

Internal change

2022-08-18 11:39:07 -07:00

python_callback_test.py

jax.test_util: add capture_stdout context manager

2022-09-12 15:21:52 -07:00

qdwh_test.py

un-skip polar/qdwh decomp tests skipped on gpu in ad6ce74

2022-08-15 12:31:43 -07:00

random_test.py

Merge pull request #12170 from froystig:just-dtype

2022-08-31 18:36:47 -07:00

remote_transfer_test.py

Increase minimum jaxlib version to 0.3.7.

2022-04-18 08:09:50 -04:00

scipy_fft_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

scipy_interpolate_test.py

Merge pull request #8625 from Edenhofer:regular_grid_interpolator

2022-01-24 12:02:34 -08:00

scipy_ndimage_test.py

[x64] make scipy_ndimage_test compatible with strict dtype promotion

2022-06-16 15:45:15 -07:00

scipy_optimize_test.py

jax.scipy.optimize: fix type inconsistency

2022-06-29 09:47:53 -07:00

scipy_signal_test.py

Remove _ prefix from functions in jax._src.dtypes.

2022-08-12 12:51:09 +00:00

scipy_stats_test.py

Roll forward: Upgrade logistic into a primitive.

2022-09-07 06:06:56 -07:00

sparse_test.py

[sparse] Lower batch-mode bcoo_dot_genernal to cusparseSpMM.

2022-09-12 10:09:41 -07:00

sparsify_test.py

[sparse] add sparse support for dynamic_slice

2022-09-01 13:42:02 -07:00

stack_test.py

Make QDWH-eig implementation jit-table.

2022-05-17 18:12:48 -07:00

state_test.py

Split State effect into Read/Write/Accum effects and tie them to Ref avals

2022-09-08 08:04:13 -07:00

stax_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

svd_test.py

[x64] make linalg functions & tests compatible with strict dtype promotion

2022-06-16 10:32:20 -07:00

transfer_guard_test.py

Bump minimum jaxlib version to 0.3.2 and remove transfer guard compatibility code

2022-04-11 15:33:27 +00:00

tree_util_test.py

Make PyTreeDef pickleable

2022-08-04 07:13:46 -07:00

util_test.py

Fix race condition for weakref destructor by catching rare exceptions.

2022-04-01 12:04:36 -07:00

version_test.py

Introduce a new jax/jaxlib versioning scheme.

2022-02-07 17:59:42 -05:00

x64_context_test.py

feat: validate jit args

2022-05-18 21:54:47 +01:00

xla_bridge_test.py

Avoid more direct references to jax._src without imports.

2022-08-24 07:51:28 -07:00

xla_interpreter_test.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

xmap_test.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00