ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

943126ca

[build] avoid top-level dependencies in jax/_src

Jake VanderPlas committed 1 year ago

e485ac95

Refactor: Add carry support to nd_loop

cperivol committed 1 year ago

718dba9b

Allow fusing all stride == 1 to support jax.lax.slice_in_dim in @fuser.fuse

Google-ML-Automation committed 1 year ago

a77cbc83

before iterating, copy dict that can be popped from via ARC / object deletion

mattjj committed 1 year ago

0d3feb8e

Merge pull request #30464 from mattjj:psend-precv-docs

Google-ML-Automation committed 1 year ago

ba3bfcae

add psend/precv to auto-generated docs

mattjj committed 1 year ago

54c618c9

Merge pull request #30447 from justinjfu:mgpu_docs_fixes

Google-ML-Automation committed 1 year ago

b32721d4

[Mosaic GPU] Fix some typos in docs

justinjfu committed 1 year ago

98665a7f

Merge pull request #30420 from melissawm:refactor-httij

Google-ML-Automation committed 1 year ago

535829d7

Merge pull request #30446 from jakevdp:source-info-util-build

Google-ML-Automation committed 1 year ago

9b47a902

Cache `canonicalize_sharding` and remove the canonicalization in `_scatter_update` since the out_sharding passed to scatter_update should already be

yashk2810 committed 1 year ago

768a64eb

avoid top-level imports in source_info_util

jakevdp committed 1 year ago

338b4ebc

Test B200 against CUDA 12.8 only

nitins17 committed 1 year ago

fd8a79f8

Add Hermetic C++ Toolchains for Linux x86_64 builds.

Google-ML-Automation committed 1 year ago

bf25960c

Merge pull request #30257 from andportnoy:aportnoy/mosaic-gpu-query-shared-memory

Google-ML-Automation committed 1 year ago

b2bd5e77

[jax:ffi] Expose `build_ffi_lowering_function` to give user access to the lowering function that returns an Op, not just Op results.

danielsuo committed 1 year ago

e381a868

[Pallas TPU] Remove `use_sreg_for_state` flag

Google-ML-Automation committed 1 year ago

4785804f

Remove out_sharding from scatter_add because we now have a scatter sharding rule

yashk2810 committed 1 year ago

1c0f24b2

Merge pull request #30440 from jax-ml:direct-linearize-context-bug

Google-ML-Automation committed 1 year ago

d5898929

Use the correct tracing context when creating zeros in linearize_from_jvp

dougalm committed 1 year ago

88eeb4fb

Implement explicit mode sharding rule for scatter

oliverdutton-iso committed 1 year ago

23563247

#sdy Add unit tests making sure we fallback to GSPMD if the original model was exported with GSPMD but loaded with Shardy.

bartchr808 committed 1 year ago

cb10a759

[jax] Export ir_attribute in jax.interpreters.mlir

James Molloy committed 1 year ago

96c7d61e

Fix typo in doc.

allanrenucci committed 1 year ago

41afeac1

[Mosaic] NFC: Just use ExtFOp and TruncFOp for bf16.

WindQAQ committed 1 year ago

9230c676

[Mosaic GPU] Simplify type constraints in MGPU dialect.

allanrenucci committed 1 year ago

8bb6b00f

[Mosaic GPU] Add support for warp shuffles with elements wider than 32-bit

apaszke committed 1 year ago

844a644b

Fix a bug where doing a CPU jax.Array -> TPU jax.Array transfer was not setting the layout correctly on the output TPU array.

yashk2810 committed 1 year ago

b57f0773

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7a6f1dfe

Older