Commits ROCm/jax - SemanticDiff

Skip test_batch_axis_sharding_jvp because of hipSPARSE issue

AratiGanesh committed 259 days ago

3f6c3187

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

Google-ML-Automation committed 259 days ago

7a5431c5

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

junwhanahn committed 259 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 259 days ago

7fedc0e5

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

emilyfertig committed 259 days ago

bdcc93f7

Disable pytorch_interoperability_test on GPU since it's broken

yashk2810 committed 259 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 259 days ago

0b986900

Merge pull request #34595 from samanklesaria:test_new_array_api

Google-ML-Automation committed 259 days ago

1aac7d88

[Mosaic GPU][NFC] Remove duplicate registration.

allanrenucci committed 259 days ago

ba4b18b6

Simplify psum_invariant's abstract_eval rule by disallowing positional axes.

yashk2810 committed 259 days ago

0e076992

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

allanrenucci committed 259 days ago

ccb5f92f

[Mosaic GPU] Ignore unknown attributes in custom call.

allanrenucci committed 259 days ago

8b1f7825

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

Google-ML-Automation committed 259 days ago

039e00d2

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

bchetioui committed 259 days ago

6b3c32f3

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

bchetioui committed 259 days ago

724eb029

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

superbobry committed 259 days ago

bd995dad

Removed a few unused functions

superbobry committed 259 days ago

1635c5be

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

apaszke committed 259 days ago

cb1214b7

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

allanrenucci committed 259 days ago

c45ec035

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

apaszke committed 260 days ago

09e023e4

Update array api test hash to most recent

samanklesaria committed 260 days ago

7fbcbd89

PallasLoweringDeterminismTest should only run on TPU backend

kanglant committed 260 days ago

9622105a

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

allanrenucci committed 260 days ago

30eaa17e

Add Bzlmod support for JAX.

meteorcloudy committed 260 days ago

228b61a7

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

Google-ML-Automation committed 260 days ago

ded22c03

Merge pull request #34586 from mattjj:config-tweak

Google-ML-Automation committed 260 days ago

359fa00e

tweak config for custom_vjp3: put in cache keys

mattjj committed 260 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 260 days ago

768b1d6d

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

Google-ML-Automation committed 260 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 260 days ago

6fa4e041