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

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -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

Remove device_replica_id_map from the Sharding interface because the standalone function should be more than enough to use. The major use-case of this is for checkpointing and accessing addressable_shards which accesses the standalone function makes it work.

2022-08-29 14:49:48 -07:00

batching_test.py

roll-forward #11952 , take 2

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

BUILD

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -07:00

callback_test.py

feat: refactor code using pyupgrade

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

checkify_test.py

Add weak type support to Array. Also make all api_test.py tests pass with Array. I have disabled the float0 test for now until I investigate.

2022-08-17 12:25:49 -07:00

clear_backends_test.py

Introduce jax.experimental.clear_backends to delete all JAX runtime backends.

2022-07-20 15:10:27 -07:00

compilation_cache_test.py

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

2022-08-23 09:36:47 -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

feat: validate jit args

2022-05-18 21:54:47 +01: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

Use .device() to get the device and platform from the device and fix TODO to point to github issue

2022-08-19 13:14:13 -07:00

dtypes_test.py

Add array test counterparts to dtypes_test.

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

dynamic_api_test.py

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

2022-08-23 09:36:47 -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

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

Rollback of upgrade logistic (sigmoid) function into a lax primitive.

2022-08-30 15:30:43 -07:00

lax_autodiff_test.py

Rollback of upgrade logistic (sigmoid) function into a lax primitive.

2022-08-30 15:30:43 -07:00

lax_control_flow_test.py

Add batching rules for state primitives and for_loop

2022-08-29 11:40:09 -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

Re-enable lax_numpy test that triggered a nonterminating LLVM compilation.

2022-08-24 11:33:30 -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

Rollback of upgrade logistic (sigmoid) function into a lax primitive.

2022-08-30 15:30:43 -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

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -07:00

multibackend_test.py

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -07:00

multiprocess_gpu_test.py

Enable one gpu per process in multinode GPU CI

2022-08-29 09:00:19 -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

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

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

pjit_test.py

Add support for MeshPspecSharding local_sharded_result_handler because SDA outputs from pjit can produce a MeshPspecSharding.

2022-08-25 17:14:05 -07:00

pmap_test.py

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -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

Enable jax_array on CPU for the entire JAX test suite!

2022-08-19 10:04:35 -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

Make jit == pjit. This means that the lowering and execution paths of jit and pjit are merged.

2022-08-29 22:03:21 -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

Rollback of upgrade logistic (sigmoid) function into a lax primitive.

2022-08-30 15:30:43 -07:00

sparse_test.py

[sparse] add support for bcoo equivalent of lax.slice

2022-08-25 09:19:44 -07:00

sparsify_test.py

[sparse] fix unary operations in presence of duplicate indices

2022-07-07 13:49:50 -07:00

stack_test.py

Make QDWH-eig implementation jit-table.

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

state_test.py

Use cases_from_list to limit number of tests

2022-08-29 14:52:14 -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