ROCm/jax

Pull Requests Commits

Merge pull request #23086 from ROCm:ci_build_fix

a-googler committed 1 year ago

3d942ef2

Merge pull request #23080 from jakevdp:array-doc

a-googler committed 1 year ago

1127f49d

Implement initial vmap over pallas_call w/ ragged inputs (via jumbles)

Google-ML-Automation committed 1 year ago

24394a1b

document jax.Array methods and attributes

jakevdp committed 1 year ago

e9d6fd37

Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.

dfm committed 1 year ago

b6306e39

Pass bazel options to requirements_update and requirements_nightly_update commands

vfdev-5 committed 1 year ago

12e8bf45

Determine LAPACK workspace during QR Factorization Kernel runtime

Paweł Paruzel committed 1 year ago

acacf888

[Easy] Refactor ragged_dot transpose, combine ragged_to_dense

Google-ML-Automation committed 1 year ago

9785368c

[ROCm] Fix python in rocm ci_build script.

Ruturaj4 committed 1 year ago

fd7c52d2

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

417fcd57

Set Clang as the default compiler in the build script.

Google-ML-Automation committed 1 year ago

a498c1e6

Merge pull request #23039 from froystig:docs3

a-googler committed 1 year ago

88ed5797

Merge pull request #23069 from jakevdp:doc-vjp

a-googler committed 1 year ago

527a4b88

Reverts 6fc57c0eb6f06b2da20c94f5f127fe4a551bda09

hawkinsp committed 1 year ago

1516d597

Rollback the change "Import from ``mlir.dialects`` lazily"

lionelfeng committed 1 year ago

322d0c2f

If `AbstractMesh` is an input to `shard_map`, then in eager mode require atleast one input to be a `NamedSharding` not all inputs.

yashk2810 committed 1 year ago

6913551d

[Pallas] Fix boolean vector loads with indexing.

justinjfu committed 1 year ago

82d3cfb3

[Pallas] Add run_scoped interpret mode rule + Enable DMA tests.

justinjfu committed 1 year ago

a18561df

Document methods of custom_jvp/custom_vjp

jakevdp committed 1 year ago

91f55129

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

2737a735

[Mosaic] Update serde to handle upstream MLIR changes

tlongeri committed 1 year ago

020513f3

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

85fb66a2

Merge pull request #23060 from jakevdp:core-deps

a-googler committed 1 year ago

acf4b324

Merge pull request #23059 from jakevdp:exports

a-googler committed 1 year ago

25cd9ea9

[Pallas] Add boolean vector support.

justinjfu committed 1 year ago

db000459

Introduce hermetic CUDA in Google ML projects.

Google-ML-Automation committed 1 year ago

599c13aa

Deprecate several internal utilities in jax.core

jakevdp committed 1 year ago

bd9698ec

Merge pull request #23040 from jakevdp:isin-method

a-googler committed 1 year ago

87a5591d

Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.

dfm committed 1 year ago

ad1bd387

Reverts 734ebd570891ceaf8c7104e12256a1edfe942b14

Google-ML-Automation committed 1 year ago

bab70dda

Newer Older