rocm_jax

benchmarks

…

filecheck

Change jax to lower the asin and atan primitives to their corresponding chlo

2022-08-10 13:05:29 -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

prepare to switch to new remat

2022-08-04 12:54:22 -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

Disable NumPy dlpack tests on GPU.

2022-08-10 09:35:41 -07:00

array_test.py

Make _device_assignment and _addressable_device_assignment a cached property because having a lru_cache as a decorator on the method with only self leaks memory. See the document below. In this case, its okay to make them a property rather than a method. Its also consistent with device_set.

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

batching_test.py

[x64] make batching_test compatible with strict dtype promotion

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

BUILD

Pmap should output SDA like Arrays to maintain the current behavior exactly. Split the shard_arg_handler for Array based on whether the mode is pmap or pjit. Why do this? The doc below explains more about the context.

2022-08-10 20:11:37 -07:00

callback_test.py

feat: refactor code using pyupgrade

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

checkify_test.py

Checkify: support checkify-of-pjit.

2022-07-29 19:25:22 +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

Internal testing change

2022-07-27 10:19:18 -07: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

add test for custom_linear_solve + new remat

2022-07-28 21:37:23 -07: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

Try flattening the locals/globals dict in the debugger and have a

2022-08-09 12:31:33 -07:00

debugging_primitives_test.py

Error in debug print if unused arguments are provided

2022-08-05 17:24:04 -07:00

distributed_multinode_test.py

Squashed commit of the following:

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

distributed_test.py

Fix MultiProcessGpuTest test

2022-07-19 11:13:32 -05:00

dtypes_test.py

Correct spelling on word

2022-07-20 18:57:12 +09:00

dynamic_api_test.py

small tweaks for bint ad

2022-08-05 08:04:50 -07:00

errors_test.py

feat: refactor code using pyupgrade

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

fft_test.py

Increase minimum NumPy version to 1.20.

2022-08-06 14:51:14 +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

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

fix ad_checkpoint.checkpoint caching issue

2022-07-29 19:59:28 -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

make tests compatible with recent pillow versions

2022-07-22 13:09:52 -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

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

Enable debug callbacks in checkpoint

2022-08-05 10:59:32 -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

Avoid top-level aliases of jax.tree_util.*

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

lax_autodiff_test.py

dynamic_slice: correctly handle negative start indices in autodiff

2022-07-21 13:41:00 -07:00

lax_control_flow_test.py

Refactor state out of for_loop

2022-08-01 15:26:55 -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

Tests: reenable some ufunc input tests

2022-08-10 12:45:32 -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

defer to custom eltype for slice lowering rule

2022-08-09 19:13:34 -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

jnp.linalg.matrix_rank: support stacks of matrices

2022-08-08 11:16:34 -07: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

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

Ensure PartitionSpec is picklable.

2022-08-08 16:56:45 -07:00

pjit_test.py

Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.

2022-08-08 19:41:42 -07:00

pmap_test.py

Pmap should output SDA like Arrays to maintain the current behavior exactly. Split the shard_arg_handler for Array based on whether the mode is pmap or pjit. Why do this? The doc below explains more about the context.

2022-08-10 20:11:37 -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

python_callback_test.py

Enable receives in TPU callbacks and add tests

2022-08-08 11:42:16 -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

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

state_test.py

Refactor state out of for_loop

2022-08-01 15:26:55 -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

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

add custom-policy partial eval and dce rules for pmap

2022-07-28 21:13:25 -07:00