ROCm/jax

Pull Requests Commits

Update rules_ml_toolchain as well as xla refs

alekstheod committed 1 hour ago

d0faeb6a

[Pallas][GPU kernel interpreter] Support barriers of multi-dimensional shape, i.e. where `num_barriers` is a tuple.

Google-ML-Automation committed 4 hours ago

3d946c66

Regenerate type stubs for nanobind v2.13.0.

Google-ML-Automation committed 4 hours ago

15db5aa6

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

Google-ML-Automation committed 6 hours ago

667f14bd

PR #38762: [Pallas][Triton] Added gpu_info and removed gpu device fallback from pallas_call_lowering

vfdev-5 committed 17 hours ago

9a428586

Expose optimization levels to Pallas

vsytch committed 20 hours ago

e6299e2e

Avoiding copy in "astype"

alvarosg committed 20 hours ago

38a697a6

[Pallas:MGPU] Add lowering for `lax.concatenate` under WG semantics.

allanrenucci committed 1 day ago

9c715422

[Mosaic SC] Change the default tiling for as_tpu_kernel to COMPACT

apaszke committed 1 day ago

a98e7468

[pallas] `mpmd_map` no longer aliases by default

superbobry committed 1 day ago

d8f8aa7c

[Pallas][GPU kernel interpreter] Support cluster barriers.

Google-ML-Automation committed 1 day ago

1946a44e

Undo XLA version bump

apaszke committed 1 day ago

5d154832

[Mosaic TPU] Bump the forward-compat version

apaszke committed 1 day ago

6e8d8e4f

[pallas:mosaic_gpu] Deprecated the `idx` parameter in `plgpu.load`

superbobry committed 1 day ago

d9aff6f7

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

Google-ML-Automation committed 1 day ago

6613f918

Make `unreduced` transfers via device_put go via `batched_copy_array_to_devices_with_sharding` when there is a single process and no reshards.

yashk2810 committed 1 day ago

956f32e2

Default core_id_mutability to False when core_ids is specified.

Google-ML-Automation committed 1 day ago

ca9d73c4

PR #38853: [mutable-arrays] fix error when returning a ref created by empty_ref

mattjj committed 1 day ago

5bb612ea

Add the constraints in `core_call_lowering` instead of `call_lowering` which is used by shard_map too. Fixes: https://github.com/jax-ml/jax/issues/38845

yashk2810 committed 1 day ago

38b88d70

hijax fixes (v2)

Cristian Garcia committed 1 day ago

8c1616e5

Reverts faf677af0abec7ea25c7a33831fcabf30a795796

yashk2810 committed 1 day ago

3a99fe37

Fix a bug where sort_jvp was not creating the iota primal with the correct sharding

yashk2810 committed 1 day ago

e50b7439

[pallas:mosaic] Use `jtu.capture_stdout` instead of `contextlib.redirect_stdout`

superbobry committed 2 days ago

4150dab4

[pallas:mosaic] Attach Pallas-level parameter names to the `FuncOp` parameters

superbobry committed 2 days ago

aee2251e

[Mosaic GPU] Mark some Mosaic GPU ops as Pure.

allanrenucci committed 2 days ago

a2b02d44

[pallas:mosaic] Re-added a few `pl.debug_print` test cases

superbobry committed 2 days ago

98edf049

[Pallas TPU] Nicer error message when a prefetched ref is not passed in explicitly

apaszke committed 2 days ago

79f101ce

[XLA:GPU]: Reattempt wrap all collectives in async-start for GPUs

sohaibiftikhar committed 2 days ago

e8f6a618

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

Google-ML-Automation committed 2 days ago

ab43b6e5

Integrate LLVM at llvm/llvm-project@f3a69752ac18

boomanaiden154 committed 2 days ago

7e2b1b5d

Older