ROCm/jax

Pull Requests Commits

Add `device_assignment` to `compile` to allow users to trace and lower with AbstractMesh and then provide devices during `compile` when it's actually needed.

yashk2810 committed 196 days ago

27a19b5c

fix ra2a vmap (batching) rule

mattjj committed 196 days ago

21a113be

[Pallas TPU] NFC: Add initial scaffolding for supporting multiple buffering, and add unit-test like test suite for basic pipelining correctness.

justinjfu committed 196 days ago

1b5e42e4

Add a field to plugin_attributes to indicate whether the PjRt plugin supports cross-host device transfers.

emilyfertig committed 196 days ago

f3338be0

[NFC][Mosaic:TPU] Minor refactor of changeTiling to make chaining packs/unpacks and relayouts easier

tlongeri committed 196 days ago

85193bdc

Include `ad_checkpoint` primitives in `roofline`.

zacmustin committed 196 days ago

45008fab

Expose CompiledMemoryStats::peak_memory_in_bytes to Python.

matthiaskramm committed 196 days ago

6b98034e

[JAX] Catch a Python exception from `nb::hash` call

hyeontaek committed 196 days ago

1e0b8cb8

[Mosaic:TPU][Relayouts] Express some retilings as unpack+pack and remove offset restrictions

tlongeri committed 196 days ago

0c59a11c

[Pallas][Mosaic GPU] Add support for TMEM Ref aliasing.

justinjfu committed 196 days ago

0b77d7c6

[Mosaic] Add tpu.enqueue_indirect_dma and verification tests

naummo committed 196 days ago

9fc8773c

Add an inline= argument to mlir.register_lowering.

hawkinsp committed 196 days ago

0fd2074c

[Mosaic] Add vmem_shared

naummo committed 196 days ago

cd72df6f

[jaxlib] Guard new transfer library API calls on `JAX_IFRT_VERSION_NUMBER`.

danielsuo committed 196 days ago

d62e385c

[JAX][Pallas] Fix pallas memory space test. Update field `output_memory_space_colors` to `output_memory_colors` in the tests.

subhankarshah committed 196 days ago

9362633b

[JAX] Remove jax.interpreters.mlir.flatten_lowering_ir_args.

hawkinsp committed 196 days ago

0b52eef0

[Mosaic GPU] Add preliminary support for scaled tcgen05.mma

apaszke committed 196 days ago

26a64c45

Update `third_party/xla/workspace.bzl` references to `third_party/xla/revision.bzl`.

danielsuo committed 196 days ago

a95976ff

Remove `jax.experimental.shard` module since all APIs have graduated to `jax.sharding`. Use `jax.sharding.reshard`, `jax.sharding.auto_axes` and `jax.sharding.explicit_axes` instead of their experimental endpoints.

yashk2810 committed 196 days ago

9bdc04c8

Migrate to `jax.sharding` endpoint for `reshard` and `auto_axes` instead of the experimental endpoint.

yashk2810 committed 196 days ago

0d2faf58

Update XLA dependency to use revision

Google-ML-Automation committed 196 days ago

8c6c437d

Merge pull request #29739 from olupton:deprecated-cufft-enums

Google-ML-Automation committed 197 days ago

7174b76c

cuFFT: remove deprecated enum values

olupton committed 197 days ago

39e03eef

[jaxpr consts] Propagate closed-over constants in lax.composite

gnecula committed 197 days ago

f419d454

[pallas:mosaic] Added `device_id` and `core_id` to `tpu.dma_wait2` op

superbobry committed 197 days ago

064c1ce5

[Mosaic:TPU] 16-bit tpu.iota

tlongeri committed 197 days ago

bf618a83

Removed unused matplotlib dependency from //tests/mosaic:gpu_test

superbobry committed 197 days ago

f54f6801

Allow all shardings if `exported.nr_devices` is 1 in `_export.py`.

ZixuanJiang committed 197 days ago

67bf78a6

Update socket transfers to call new transfer library APIs.

emilyfertig committed 197 days ago

877bee73

Update ml_dtypes to 0.5.1 to align with JAX and TensorFlow

apivovarov committed 197 days ago

21b52fb5

Newer Older