Commits ROCm/jax - SemanticDiff

Skip test_batch_axis_sharding_jvp because of hipSPARSE issue

AratiGanesh committed 3 days ago

3f6c3187

Do not cache executables when MeshComputations.compile is called with a device list.

Google-ML-Automation committed 3 days ago

7a5431c5

Make compile methods in `xla::ifrt::Compiler` return a future of an executable

junwhanahn committed 3 days ago

3d1ea0b2

Introduce `jax.extend.clear_in_memory_compilation_cache` to allow users to clear the LRU cache for mesh executables during compilation

krishnaharidasan committed 3 days ago

7fedc0e5

Support cross-host `device_put` for a few new cases.

emilyfertig committed 3 days ago

bdcc93f7

Disable pytorch_interoperability_test on GPU since it's broken

yashk2810 committed 4 days ago

3cf7ece1

Fix error checking logic in check_unreduced_args to only error if axes of a primitive and aval's unreduced intersect.

yashk2810 committed 4 days ago

0b986900

Merge pull request #34595 from samanklesaria:test_new_array_api

Google-ML-Automation committed 4 days ago

1aac7d88

[Mosaic GPU][NFC] Remove duplicate registration.

allanrenucci committed 4 days ago

ba4b18b6

Simplify psum_invariant's abstract_eval rule by disallowing positional axes.

yashk2810 committed 4 days ago

0e076992

[Mosaic GPU][NFC] Add null checks before nullable pointers are dereferenced.

allanrenucci committed 4 days ago

ccb5f92f

[Mosaic GPU] Ignore unknown attributes in custom call.

allanrenucci committed 4 days ago

8b1f7825

[Pallas:MGPU] Lower multi-device operation with collective metadata.

Google-ML-Automation committed 4 days ago

039e00d2

[Mosaic GPU] Override the PTX version used LLVM->PTX compilation with the highest available one during PTX->SASS compilation.

bchetioui committed 4 days ago

6b3c32f3

[Pallas/Mosaic GPU] Add the `reduction_scratch_bytes` field to `CompilerParams`.

bchetioui committed 4 days ago

724eb029

[mosaic] `as_tpu_kernel` now accepts `tiling=`

superbobry committed 4 days ago

bd995dad

Removed a few unused functions

superbobry committed 4 days ago

1635c5be

[Pallas:TPU] Add a missing libTPU version check for explicit MXU API tests

apaszke committed 4 days ago

cb1214b7

[Mosaic GPU] Initialise NVSHMEM runtime library as part of `InitKernel`.

allanrenucci committed 4 days ago

c45ec035

[Mosaic GPU] Add support for all sensible conversions involving f8

apaszke committed 4 days ago

09e023e4

Update array api test hash to most recent

samanklesaria committed 4 days ago

7fbcbd89

PallasLoweringDeterminismTest should only run on TPU backend

kanglant committed 4 days ago

9622105a

[Mosaic GPU] Use existing CUDA compute capability whenever available.

allanrenucci committed 4 days ago

30eaa17e

Add Bzlmod support for JAX.

meteorcloudy committed 4 days ago

228b61a7

Update XLA dependency to use revision http://github.com/openxla/xla/commit/a9cb3c90c30895103ffbb4fc231b9930d5c61b91

Google-ML-Automation committed 4 days ago

ded22c03

Merge pull request #34586 from mattjj:config-tweak

Google-ML-Automation committed 4 days ago

359fa00e

tweak config for custom_vjp3: put in cache keys

mattjj committed 4 days ago

d87942f3

Check for args_not_mapped first and then fancy_primitive_batchers. If not, it leads to numerical issues (no idea why, but this was the behavior before skippable_batchers was deleted)

yashk2810 committed 4 days ago

768b1d6d

Merge pull request #34573 from jakevdp:refactor-hijax-tests

Google-ML-Automation committed 4 days ago

46d36125

Move the pallas_call batching logic which removes explicit_mesh_axis_names from AxisEnv into core.py and enter into that context in BatchTrace.process_primitive

yashk2810 committed 4 days ago

6fa4e041