ROCm/jax

Pull Requests Commits

Update rules_ml_toolchain to integrate new CUPTI headers.

ybaturina committed 77 days ago

12553810

[pyrefly] type-check jax.experimental/array_serialization & colocated_python

jakevdp committed 77 days ago

180b30f3

Add constraints on output during cond, while_loop, sort lowering

yashk2810 committed 77 days ago

877ca781

shard map without linear util

dougalm committed 77 days ago

e58a022a

Merge pull request #35775 from jakevdp:pyrefly-experimental

Google-ML-Automation committed 77 days ago

0936349f

[Mosaic TPU] Strengthen the tpu.memref_slice verifier

apaszke committed 77 days ago

fc635497

Increase size of `ode_test`.

mwhittaker committed 77 days ago

ef13248c

[pallas:mosaic] Fixed the signature of `infer_tiling` -- it never returns `None`

superbobry committed 77 days ago

bcb3e200

[NFC] Extract reduce expression reduction logic into a helper function.

allanrenucci committed 77 days ago

fc2e5146

[Mosaic GPU] Remove workaround for miscompilation

apaszke committed 77 days ago

c1f4c33a

[Mosaic GPU] Add support for atomic stores from arrays with the WG strided layout

apaszke committed 77 days ago

42c4d0a3

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

Google-ML-Automation committed 77 days ago

4b578528

[Mosaic GPU] Clarify semantic of `BroadcastInDimOp`.

allanrenucci committed 77 days ago

336ca0c8

[Pallas:MGPU] Add a test for `broadcast_in_dim` of mixed size-1 and size-0 dimensions.

allanrenucci committed 77 days ago

3071c638

[NFC] Remove redundant `num_threads` and `thread_name` parameters in tests.

allanrenucci committed 77 days ago

2baa52ba

[Mosaic TPU] Specialize indices of grid dims of size 1 to static 0

apaszke committed 77 days ago

ac71172e

[Pallas:TPU] Drop the buffer count to 1 when the pltpu.emit_pipeline operand window is constant

apaszke committed 77 days ago

1d2ef4cc

[Mosaic GPU] Add support for keeping dimensions in `Reduce` constraint.

allanrenucci committed 77 days ago

53346b39

[Mosaic GPU] Add bitwidth to IsTransferable.

golechwierowicz committed 78 days ago

b94ce497

[Pallas:MGPU] Expose the leader_tracked in copy_gmem_to_smem

apaszke committed 78 days ago

d8bbe601

Merge pull request #35796 from gnecula:ser_abs_dev2

Google-ML-Automation committed 78 days ago

422cf9d2

[export] Add support for serializating AbstractMesh.abstract_device

gnecula committed 78 days ago

e812c59c

[Pallas] Add excluding_{lhs,rhs}_dims to jax_dot_dims_to_tpu_dot_dot_dims.

WindQAQ committed 78 days ago

2acb4492

Merge pull request #35793 from mattjj:is-vjp

Google-ML-Automation committed 78 days ago

d7ccdd1e

[JAX] Require all positional arguments passed to bind to be JAX values.

hawkinsp committed 78 days ago

c43f0a0f

extremely complex fix

mattjj committed 78 days ago

04581201

jaxlib: Also let the message attribute be optional for JaxRuntimeError.

pizzud committed 78 days ago

284116c7

Remove references to has_cpu_sparse, which is always true.

hawkinsp committed 78 days ago

5a9490fa

[Mosaic] Generate enums in Python bindings

tlongeri committed 78 days ago

ebf39ff6

jaxlib: Expose the status attributes on XlaRuntimeError.

pizzud committed 78 days ago

ffcc5ea9

Newer Older