mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/tests
History
jax authors 7e372944f9 Fix the missing cache_misses metric when min compile time is set to zero. 
Remove the code which checks if the min compile time is greater than zero. After this change, we can catch cache_misses when min compile time is zero.

Testing: revised unit test.
PiperOrigin-RevId: 579951415
2023-11-06 14:04:35 -08: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
Remove temporary patch for where broadcasting test
2023-10-23 10:37:25 -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
In api_test.py, wait for the result in test_double_donation.
2023-11-03 12:23:55 -07: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
[IFRT] Fix incorrect type numbers for e4m3 and e5m2 types.
2023-10-24 14:38:00 -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
Deprecated define_* and DEFINE_* methods on jax.config
2023-10-29 20:58:19 +00:00
cache_key_test.py
Fall-back to original device/backend hashing if topology-desc is unavailable.
2023-11-02 18:43:48 -07: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
Fix the missing cache_misses metric when min compile time is set to zero.
2023-11-06 14:04:35 -08:00
config_test.py
Deprecated define_* and DEFINE_* methods on jax.config
2023-10-29 20:58:19 +00: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
MAINT Do not import the config object in JAX internals
2023-10-18 10:55:13 +01:00
export_test.py
Add a lax.platform_dependent API for writing platform-dependent code.
2023-11-02 14:31:38 +01:00
extend_test.py
[jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl
2023-10-31 17:21:54 -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
Disable flaky python callback test.
2023-10-23 12:24:05 -07: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
Add a lax.platform_dependent API for writing platform-dependent code.
2023-11-02 14:31:38 +01: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
Relax some test tolerances for TPU.
2023-10-24 10:45:40 -07: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
jnp.cross: account for numpy 2.0 deprecation
2023-11-03 14:15:23 -07: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
Add beta function
2023-11-05 15:37:38 -08: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
[JAX] Add an option subset_by_index that allows computing a contiguous subset of eigenvalues from eigh.
2023-10-27 09:29:33 -07: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
Places the remat decorator on top of the body function.
2023-10-27 15:27:19 -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
MAINT Do not import the config object in JAX internals
2023-10-18 10:55:13 +01: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
Relax some test tolerances for TPU.
2023-10-24 10:45:40 -07: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
[Try again] For nested pjit's cache the generation of StableHLO if it satifies the key. This should help in improving the lowering time.
2023-10-30 15:28:43 -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
Enable python_callback_test for stream executor.
2023-11-02 13:26:59 -07: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
Deprecate jax.random.shuffle
2023-11-06 12:21:56 -08:00
random_test.py
[random] better errors for unsupported operations on prng keys
2023-11-03 19:23:18 -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
MAINT Do not import the config object in JAX internals
2023-10-18 10:55:13 +01:00
scipy_stats_test.py
Merge pull request #18126 from niqodea:wrapcauchy
2023-10-18 13:18:20 -07:00
shard_map_test.py
output res forwarding optimization for shard_map and jit
2023-10-18 23:56:26 -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
[XLA:GPU] Consider Triton for all non-pure GEMM fusions
2023-11-04 16:05:19 -07:00