ROCm/jax

Pull Requests Commits

Merge pull request #33021 from cgarciae:hijax-coverage-tests

Google-ML-Automation committed 206 days ago

f386951c

Update XLA dependency to use revision http://github.com/openxla/xla/commit/904baf4249be108fd349db443a313294dbbb7f69

Google-ML-Automation committed 206 days ago

b11d86c8

[jax:ffi] Declare ffi::TypeInfo as a struct static member

ezhulenev committed 207 days ago

4f88b028

[Mosaic TPU] Allow small sublane tilings in second/third minor transpose for unpacked dtype instead of forcing native tiling, which can be more efficient when the second minor dimension of input tensor is small.

yueshengys committed 207 days ago

af44ec5c

[xla:ffi] Migrate to xla::ffi::MakeTypeInfo() API

ezhulenev committed 207 days ago

c5c103cd

Reverts f72308e83aabcb061b16d58fc8330587de3cee02

hawkinsp committed 207 days ago

437d260f

add grad hijax coverage tests

cgarciae committed 207 days ago

709e3364

Merge pull request #33041 from mattjj:hijax-vjp-accums

Google-ML-Automation committed 207 days ago

f6edb899

Pass git hash to artifact build for release and nightly builds

nitins17 committed 207 days ago

40e9dc9d

[hijax] upgrade VJPHiPrimitive.def_bwd to work with accums

mattjj committed 207 days ago

eff56af4

[CI] Reduce the number of concurrent test jobs in the TSAN workflow.

belitskiy committed 207 days ago

08c63d38

[Mosaic GPU][NFC] Replace calls to `llvm.getelementptr` with `utils.getelementptr`.

allanrenucci committed 207 days ago

0a972ed4

[Mosaic GPU] Remove `TraversalOrder` from `traverse_op`.

bchetioui committed 207 days ago

7239d15f

[Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,*},(8,128) -> {0,*},(4,128) has been just supported.

yueshengys committed 207 days ago

dc7d56e6

[Pallas/Mosaic GPU] Add support for `s8` WGMMA with lhs in registers.

bchetioui committed 207 days ago

7a007ea9

[Mosaic GPU][NFC] Silence some linter warning.

bchetioui committed 207 days ago

14d89b51

Repurpose `tpu_sparsecore` annotation for offloading purposes

vsytch committed 207 days ago

2d1479dc

[Mosaic GPU] Add partial support for `s8` lhs in registers in `wgmma`.

bchetioui committed 207 days ago

afbd5ff1

[doc] specify help string by keyword in configs.

jakevdp committed 207 days ago

fbb7be6a

Reverts b086ca2b1200f99b4e93e17513bae0bf57158d5d

pravnar committed 207 days ago

3ff628cb

[Pallas:SC] Properly select the serialized IR version to avoid libtpu incompatibility

apaszke committed 207 days ago

2c23d5df

Merge pull request #33019 from emilyfertig:emilyaf-fix-extending-pytrees

Google-ML-Automation committed 207 days ago

aaa2ce8e

Fix extending-pytrees link.

emilyfertig committed 207 days ago

c8cfdb65

Disable test_streamed_gemm_overlap on jaxlib 0.8.0.

hawkinsp committed 207 days ago

e2345fb4

[Mosaic GPU] Bump tolerance in the matmul test

apaszke committed 207 days ago

39288614

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

Google-ML-Automation committed 207 days ago

f6542ab3

Cache MetaTy creation and special case `Tracer` because `getattr(Tracer, 'sharding', None)` is expensive.

yashk2810 committed 208 days ago

37983f13

[hijax] rename NewstyleHiPrimitive, expose in jax.experimental

mattjj committed 208 days ago

0e5e5c75

Reverts 03614c96254014c735e234aec1258d448d62e354

Google-ML-Automation committed 208 days ago

f72308e8

[Pallas] Implement `delay` on MGPU and move `pltpu.delay` to the `pl` namespace.

justinjfu committed 208 days ago

91a6f70d

Newer Older