rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-25 02:36:06 +00:00

History

George Necula 82a2793fc9 [export] Improve the calling of multi-platform exported module

Previously we declared the lowering rule for call_exported to be
platform specific. This was correct, but in the case when the
caller function is lowered itself for multiple platforms this results
in multiple copies of the inner called Exported. Now instead we
make the call_exported rule be platform independent and make it
compute the platform index for the called module based on the
platform index in the caller module. This results in a single
copy of the HLO for the called module in the output.

2023-10-19 17:40:46 +02:00

filecheck

Make mlir.custom_call() more general and expose it as jax.interpreters.mlir.custom_call().

2023-08-29 08:50:07 -07:00

notebooks

Delete remote TPU support.

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

pallas

[pallas] dot fixes.

2023-10-16 12:35:43 -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

Remove test code that checks for the se_tpu runtime.

2023-09-25 09:30:07 -07:00

api_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

array_test.py

Optimize make_array_from_callback for fully replicated shardings by going via batched_device_put

2023-10-10 19:02:04 -07:00

batching_test.py

Fix an A100 nightly unit test failure on testDotGeneral() by replacing TF32 with float32

2023-09-12 16:31:22 -07:00

BUILD

random_lax_test: Bump shards number for CPU config.

2023-10-17 12:59:12 -07:00

cache_key_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

checkify_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

Instrument metrics for the new JAX compilation cache key generation algorithm.

2023-10-12 14:56:02 -07:00

core_test.py

MAINT Use a generator expression with all() and any()

2023-10-10 22:33:03 +01:00

custom_linear_solve_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -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

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

debug_nans_test.py

Change to simpler import for jax.config

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

debugger_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

debugging_primitives_test.py

Return PositionalSharding if input's rank is >= 3 or a NamedSharding if a mesh is available via the context from inspect_array_sharding. Never return GSPMDSharding from inspect_array_sharding.

2023-10-12 16:55:12 -07:00

deprecation_test.py

Fix tests that mistakenly used assertRaises(..., msg=...) to match a message.

2023-10-06 08:07:49 -04:00

dtypes_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

dynamic_api_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

errors_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

experimental_rnn_test.py

Fix test breakage in RNN test with old jaxlibs.

2023-09-20 11:50:04 -04:00

export_test.py

[export] Improve the calling of multi-platform exported module

2023-10-19 17:40:46 +02:00

extend_test.py

test custom PRNG impl construction round trip

2023-10-06 11:18:03 -07:00

fft_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

for_loop_test.py

Tests: set jax_legacy_prng_key='error'

2023-08-28 10:56:09 -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

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

host_callback_to_tf_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jax_to_ir_test.py

Delete jax_experimental_name_stack flag

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

jaxpr_effects_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jaxpr_util_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jet_test.py

Tests: set jax_legacy_prng_key='error'

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

lax_autodiff_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_control_flow_test.py

Fix test failure in lax_control_flow_test on Mac ARM.

2023-10-09 10:34:55 -04:00

lax_numpy_einsum_test.py

jnp.einsum: lower to mixed-precision dot_general when possible.

2023-08-24 10:31:39 -07:00

lax_numpy_indexing_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_numpy_operators_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_numpy_reducers_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_numpy_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_numpy_ufuncs_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_scipy_special_functions_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_scipy_spectral_dac_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_scipy_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_vmap_op_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_vmap_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lobpcg_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

logging_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

memories_test.py

Add GetOpSharding to XLA/PjRt utils.

2023-10-17 15:46:52 -07:00

mesh_utils_test.py

Any devices passed to jax.sharding.Mesh are required to be hashable.

2023-08-29 12:20:54 -07:00

metadata_test.py

Change metadata_test to tolerate paths with backslashes.

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

mock_gpu_test.py

Fix mock_gpu_test on OSS build.

2023-10-03 10:55:20 -07:00

monitoring_test.py

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -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

Tests: set jax_legacy_prng_key='error'

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

multibackend_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

multiprocess_gpu_test.py

MAINT Drop underscore from the name of externally-referenced state objects

2023-10-13 21:30:13 +01:00

name_stack_test.py

Propagate ad.Zeros to the scan body function for jax.lax.scan for the output 'ys'.

2023-09-07 14:36:53 -07:00

nn_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

Re-enable the test as the GPU plugin profiler fixes are in.

2023-10-16 21:56:46 -07:00

pickle_test.py

Tests: set jax_legacy_prng_key='error'

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

pjit_test.py

[random] remove internal uses of deprecated prng.seed_with_impl()

2023-10-17 13:18:08 -07:00

pmap_test.py

[XLA:GPU] Fix bug in all-to-all for complex data types.

2023-10-16 16:02:22 -07:00

polynomial_test.py

Clean up some device opt-in/opt-outs in test suite.

2023-09-27 14:56:00 -07:00

profiler_test.py

Fix testProfilerGetFDOProfile.

2023-10-16 14:52:44 -07:00

python_callback_test.py

MAINT Drop underscore from the name of externally-referenced state objects

2023-10-13 21:30:13 +01:00

pytorch_interoperability_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

qdwh_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

random_lax_test.py

[random] remove internal uses of deprecated prng.seed_with_impl()

2023-10-17 13:18:08 -07:00

random_test.py

[random] make PRNG impl attributes private

2023-10-18 11:10:47 -07: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

Test: adjust tolerances for upstream scipy.minimize

2023-09-25 12:20:34 -07:00

scipy_signal_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

scipy_spatial_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

scipy_stats_test.py

Merge pull request #18126 from niqodea:wrapcauchy

2023-10-18 13:18:20 -07:00

shard_map_test.py

[shard_map] fix axis env extension bug

2023-10-16 12:36:51 -07:00

source_info_test.py

Change to simpler import for jax.config

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

sparse_bcoo_bcsr_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

sparse_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Drop underscore from the name of externally-referenced state objects

2023-10-13 21:30:13 +01:00

stax_test.py

Fix some numpy 2.0 incompatibilities

2023-09-21 10:24:52 -07:00

svd_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

transfer_guard_test.py

Change to simpler import for jax.config

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

tree_util_test.py

Add tests to cover PyTreeDef.flatten_up_to error scenarios.

2023-10-02 13:00:55 -07:00

typing_test.py

Fix or ignore some pytype errors.

2023-10-18 16:21:59 -07:00

util_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

version_test.py

Add more extensive tests for version strings

2023-09-08 11:33:49 -07:00

x64_context_test.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

xla_bridge_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01: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

MAINT Drop underscore from the name of externally-referenced state objects

2023-10-13 21:30:13 +01:00