rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

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

filecheck

Integrate LLVM at llvm/llvm-project@71c9757474

2022-07-11 14:21:09 -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

[dynamic-shapes] Split dynamic_api_test.py

2022-07-14 20:18:53 -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

[ROCm] Disable new array_interoperability dlpack tests.

2022-07-17 04:48:11 +00:00

array_test.py

Make astype work for Array that are sharded. The current behavior is the same as SDA i.e. it round trips via host.

2022-06-28 12:49:12 -07:00

batching_test.py

[x64] make batching_test compatible with strict dtype promotion

2022-06-17 16:08:54 -07:00

BUILD

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

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

callback_test.py

feat: refactor code using pyupgrade

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

checkify_test.py

Checkify: add (checkify-of-)vmap-of-check.

2022-06-27 10:34:26 +01: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

Run compilation_cache_test.py using TFRT

2022-02-04 10:17:59 -08:00

core_test.py

Merge pull request #11151 from mattjj:input-fwd-residual-optimization

2022-06-17 15:43:44 -07:00

custom_linear_solve_test.py

Relax test tolerances.

2022-07-03 15:13:28 -04:00

custom_object_test.py

add support for lambda-bound dynamic shape output (iree only)

2022-05-18 21:57:30 -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

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

debugging_primitives_test.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

distributed_test.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

dtypes_test.py

BUG: ensure that boolean scalars are never marked weak

2022-07-07 15:41:23 -07:00

dynamic_api_test.py

Merge pull request #11539 from gnecula:ds_reshape

2022-07-19 23:13:03 -07:00

errors_test.py

feat: refactor code using pyupgrade

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

fft_test.py

[x64] make fft functionality compatible with strict dtype promotion

2022-06-15 10:10: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

Add a util to fetch value of a GDA to host when its single-controller. Error out in McJAX

2022-07-07 11:09:13 -07:00

heap_profiler_test.py

Move jax.test_util to jax._src.test_util.

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

host_callback_test.py

Rollback of HCB GPU custom call due to internal failures

2022-07-10 13:05:27 -07:00

host_callback_to_tf_test.py

Move jax.test_util to jax._src.test_util.

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

image_test.py

feat: refactor code using pyupgrade

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

infeed_test.py

Avoid top-level aliases of jax.tree_util.*

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

jax_jit_test.py

feat: validate jit args

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

jax_to_ir_test.py

Fix name stack bugs

2022-04-19 11:14:41 -07:00

jaxpr_effects_test.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

jaxpr_util_test.py

Increase minimum jaxlib version to 0.3.0.

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

jet_test.py

Avoid top-level aliases of jax.tree_util.*

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

lax_autodiff_test.py

BUG: fix jvp rule for lax.rem

2022-07-12 09:50:42 -07:00

lax_control_flow_test.py

Add loop invariant residual fixpoint test

2022-07-11 13:10:03 -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

[x64] make lax_numpy_indexing_test pass with strict dtype promotion

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

lax_numpy_test.py

Add dlpack support to device_array and jax.numpy

2022-07-15 17:31:11 -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

[x64] make jax.scipy.sparse.linalg compatible with strict dtype promotion

2022-06-17 14:04:05 -07:00

lax_scipy_test.py

[x64] make lax_scipy_test.py compatible with strict dtype promotion

2022-06-14 10:02:45 -07:00

lax_test.py

Merge pull request #11114 from jakevdp:x64-lax-test

2022-06-15 17:44:23 -07:00

lax_vmap_test.py

[x64] make vmap_test compatible with strict dtype promotion

2022-06-17 15:55:08 -07:00

linalg_test.py

jax.scipy.linalg.solve: deprecate the sym_pos argument following scipy 1.9.0

2022-07-19 13:57:49 -07:00

lobpcg_test.py

[ROCm] Disable lobpcg unit test for ROCm until performance issue is resolved.

2022-07-18 18:12:19 +00:00

masking_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01: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

Fix name stack bugs

2022-04-19 11:14:41 -07:00

multi_device_test.py

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

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

multibackend_test.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01: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

Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.

2022-03-23 20:55:22 +00: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

DeviceArray: Improve support for copy, deepcopy, and pickle

2022-05-19 12:00:58 -07:00

pjit_test.py

OpSharding doesn't have __eq__ defined on it. Don't check sharding equality using opsharding until it does support that.

2022-07-20 15:03:39 -07:00

pmap_test.py

Merge pull request #11407 from hawkinsp:minver

2022-07-08 06:04:47 -07:00

polynomial_test.py

jnp.roots: better support for computation under JIT

2022-06-23 14:48:53 -07:00

profiler_test.py

Add collect_profile script

2022-06-03 17:56:17 -07:00

qdwh_test.py

feat: refactor code using pyupgrade

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

random_test.py

random.choice: make compatible with strict promotion

2022-06-28 10:41:30 -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

[x64] make jax.scipy.signal compatible with strict dtype promotion.

2022-06-21 09:28:46 -07:00

scipy_stats_test.py

looser TPU precision

2022-06-28 17:07:30 -04:00

sharded_jit_test.py

[x64] make sharded_jit_test compatible with strict dtype promotion

2022-06-17 16:42:45 -07:00

sparse_test.py

[sparse] make sparse objects compatible with jax.jit.lower()

2022-07-15 09:58:31 -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

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

Reduce the verbosity of treedef printing for custom nodes.

2022-07-15 07:14:28 -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

Only register backend factories once when jax_platforms config value is set.

2022-05-20 09:46:37 -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

Merge pull request #11367 from apaszke:xmap-tracer-leak

2022-07-07 02:01:51 -07:00