ROCm/jax

Pull Requests Commits

Merge pull request #32972 from mattjj:fix-newstyle-hijax-abstract-eval-caching

Google-ML-Automation committed 308 days ago

297fc6ba

[JAX] Require IFRT values and executables to have some `xla::ifrt::UserContext` if they are wrapped as JAX objects

hyeontaek committed 308 days ago

d9ba98c5

jnp: add mean argument to jnp.std and jnp.var families

arvindsankar committed 308 days ago

84954db2

[hijax] fix caching logic for newstyle hi primitives

mattjj committed 308 days ago

a4b99ac3

reshape->store into strided_store rewrite

voznesenskym committed 308 days ago

3d9f1956

Extend `isGuaranteedDivisible to pass on Add op when both sides are divisible.

IvyZX committed 308 days ago

74c6b639

[Pallas:MGPU] Add Warpgroup lowering for `lax.slice_p`.

allanrenucci committed 308 days ago

aa364711

[Mosaic GPU] Fix op traversal during layout inference to be truly pre-order.

dimitar-asenov committed 308 days ago

19e1d689

[Mosaic] Fix JSON formatting in TPU custom call config.

Google-ML-Automation committed 308 days ago

392be5b3

Use more conventional type naming in `_type_promotion_lattice`.

chr1sj0nes committed 308 days ago

9f6a1aea

[Mosaic GPU][NFC] Rename `OperandOrResult` to `ValueSite`.

dimitar-asenov committed 308 days ago

f0bf93b7

[Mosaic GPU][NFC] Add the `Argument` VariableType.

dimitar-asenov committed 308 days ago

74612654

[Mosaic GPU] Extend `Divides` constraint to support `RegisterLayout`.

allanrenucci committed 308 days ago

8889828c

Updated jaxlib stubs by running the stubgen with --exclude-values

superbobry committed 308 days ago

357e16dd

[Pallas:SC] Adds support for `VMEM_SHARED` and `HBM` in `plsc.kernel`/`pl.core_map` (with Scalar and VectorSubcore meshes) via `pl.get_global`.

brianwa84 committed 308 days ago

ac6a41f9

[Mosaic GPU] Introduces `WG_SPLAT` as a supported `RegisterLayout` in `test_util.py`.

allanrenucci committed 308 days ago

983c5925

[pmap] Disable testAllToAllSplitAxis for GPUs.

danielsuo committed 308 days ago

7cd239b0

Add a flag for the coordination service initialization timeout for multiprocess tests, increase it for tsan (roll-forward with fix).

emilyfertig committed 308 days ago

da35ee8d

[Mosaic GPU] Use `test_util.RegisterLayout` in layout inference tests.

allanrenucci committed 308 days ago

1d07a32c

Add/correct skips for failing tpu_pallas(_pipeline)?_test_tpu_7x test cases.

belitskiy committed 308 days ago

35f0ab10

Fix infinite loop bug in index_map from get_p pull block spec rule.

jburnim committed 308 days ago

11333235

[pmap] Mark `PythonPmapTest.testAllToAllSplitAxis` as thread-unsafe.

danielsuo committed 308 days ago

e3f37d43

[pallas:sc] Removed an unnecessary check from the tpu.store_vector_idx verifier

superbobry committed 308 days ago

e20c4024

[Mosaic GPU][NFC] Extract `RegisterLayout` test utility into a separate library.

allanrenucci committed 308 days ago

f3bf855e

[Mosaic TPU] Refactor and generalize the load/reshape fusion

apaszke committed 308 days ago

58bbb9a6

[pallas] Deprecated `pl.swap`

superbobry committed 308 days ago

f1a8961c

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

Google-ML-Automation committed 308 days ago

4cfd54d1

Reverts f6101c5b56427e21bee1ed202b5e99f193f87780

Google-ML-Automation committed 308 days ago

f967e4c1

Reverts cf2211c4688a72720d89ce888f2f104c2a626654

Google-ML-Automation committed 308 days ago

e6a19189

sharadmv committed 308 days ago

f6101c5b

Newer Older