ROCm/jax

Pull Requests Commits

[Pallas:MGPU] Add a collective matmul tutorial

apaszke committed 120 days ago

dfdb0dc6

[hijax] add a QArray example using new-style hijax primitives

mattjj committed 120 days ago

8fe978cd

[mutable-arrays] fix batching bugs in ref primitives

mattjj committed 120 days ago

387293f9

[Mosaic TPU] Relax second minor dim shape restrictions for packed second/third minor transpose

apaszke committed 120 days ago

b90a584e

[Mosaic GPU] Enforce `CustomPrimitiveOp` has a single-block region.

allanrenucci committed 120 days ago

6bda0bf6

Switch the default nonsymmetric eigendecomposition implementation to use cusolver.

hawkinsp committed 120 days ago

97b001a9

[Mosaic] Use small second minor tiling for packed type reshape of contiguous vreg slice.

WindQAQ committed 120 days ago

d48009f6

[Mosaic GPU] Support profiling for WG semantic.

allanrenucci committed 120 days ago

d7f9de88

[pallas:sc] Allowed scalar loads and stores to/from VMEM

superbobry committed 120 days ago

ae0ea1df

Move epilogue SMEM allocation to the kernel function

apaszke committed 120 days ago

f4448517

Automated Code Change

Google-ML-Automation committed 120 days ago

a0e49c27

[Mosaic GPU] Rewrite `CustomPrimitiveOp` lowering.

allanrenucci committed 120 days ago

d0c12f00

[mosaic] Removed the deprecated `tpu.wait_dma` op

superbobry committed 121 days ago

e2a4c1a5

[Pallas] Add pl.get_global for global semaphore allocations

apaszke committed 121 days ago

274f9bba

[Mosaic GPU][NFC] Derive `smem_buffer_ptr` in `OnDeviceProfiler.record`.

allanrenucci committed 121 days ago

8f6a72d8

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

Google-ML-Automation committed 121 days ago

53bbba06

[Mosaic TPU] Trace all basic blocks by default

Google-ML-Automation committed 121 days ago

4ee9cb2b

[Mosaic] More informative Vmem OOM error message.

yueshengys committed 121 days ago

ccf90671

[jax] Explicitly declare ones as jnp.float32

basioli-k committed 121 days ago

bc33c63d

Merge pull request #32478 from jax-ml:newstyle-hiprimitive

Google-ML-Automation committed 121 days ago

5566c346

Add an API for defining a new hijax primitive in terms of expand/fwd/bwd/batch

dougalm committed 121 days ago

4dcfaf53

Merge pull request #32491 from johannahaffner:patch-2

Google-ML-Automation committed 121 days ago

227d7e24

Remove CuDNN version checks that are always true.

hawkinsp committed 121 days ago

fd9d71c3

[Mosaic] Fix the usage of interleaved packing/unpacking to align with the semantics in tpu.td.

WindQAQ committed 121 days ago

f1c6cc49

Fix typo in Cholesky decomposition docstring

johannahaffner committed 121 days ago

Verified 2396ed4e

Avoid putting const arrays in kwargs of custom_fusion_p.bind.

jburnim committed 121 days ago

d64e7040

Small cleanups in mlir.py.

hawkinsp committed 121 days ago

c98e5a64

Disables ann_test.py on tsan TPU because of timeouts.

mwhittaker committed 121 days ago

2892bf1e

[jaxlib] Generate .pyi files for the _jax extension via nanobind's stubgen

superbobry committed 121 days ago

b53b6347

Include a sequence number when dumping jaxprs. If two functions have the same name we may overwrite jaxprs with later ones.

hawkinsp committed 121 days ago

bcbbaece

Newer Older