ROCm/jax

Pull Requests Commits

Update install docs for AMDGPU

mrodden committed 1 year ago

d52c8723

Merge pull request #28150 from dfm:gh28144

Google-ML-Automation committed 1 year ago

0e25620a

[Pallas Fuser] Allow multiple BlockSpec inputs to select_n push rule if they are identical

sharadmv committed 1 year ago

e81dae61

Merge pull request #28100 from justinjfu:pipe_docs_v2

Google-ML-Automation committed 1 year ago

c7417ce0

Fix handling of SymbolicZero output when batching custom_jvp.

dfm committed 1 year ago

82a79c59

[Pallas] Update TPU pipelining docs

justinjfu committed 1 year ago

64c645dc

[CI] Add tpu v6e-8 to nightly test and release test.

MichaelHudgins committed 1 year ago

91deb25e

[Pallas][Mosaic TPU] Improve error message for 1D block specs when the block size is too small.

justinjfu committed 1 year ago

720488fd

Merge pull request #28126 from jakevdp:scatter-annotations

Google-ML-Automation committed 1 year ago

711ba9a9

Merge pull request #24074 from scottstanie:fix-polyfit-cov

Google-ML-Automation committed 1 year ago

5b839ff3

Merge pull request #28092 from dfm:gh25847

Google-ML-Automation committed 1 year ago

437d7d88

Merge pull request #27574 from jburnim:jburnim_pallas_core_map

Google-ML-Automation committed 1 year ago

3c650290

Merge pull request #28166 from hawkinsp:lockfiles

Google-ML-Automation committed 1 year ago

9326d0bd

Merge pull request #28106 from jakevdp:frexp-grad

Google-ML-Automation committed 1 year ago

76b440c8

Fix ensure_compile_time_eval context after stackless.

dfm committed 1 year ago

6f3a8b4a

Add support for generating freethreaded pip lockfiles.

hawkinsp committed 1 year ago

c477cb27

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

59a54650

typing: improve annotations for scatter implementations

jakevdp committed 1 year ago

972cec16

[mosaic_gpu] Fixed the export code path in `_mosaic_gpu_lowering_rule`

superbobry committed 1 year ago

dd472e47

#sdy Add more debug info when there is a mesh mismatch in JAX export

bartchr808 committed 1 year ago

b12e05c5

Merge pull request #28098 from andportnoy:mosaic-gpu-profiler-padding

Google-ML-Automation committed 1 year ago

541768c6

[JAX] Use `xla::ifrt::Client::MakeArraysFromHostBufferShards()` in Array creation when possible

hyeontaek committed 1 year ago

f346fd09

[Pallas][jax] Better error message for unexpected types in standard abstract eval

Google-ML-Automation committed 1 year ago

d0b6eb2e

[ragged-paged-attn] Add auto-tuned table that is being used for vLLM

bythew3i committed 1 year ago

640de9e1

[Pallas] Support the new TPU interpret mode with pl.core_map.

jburnim committed 1 year ago

33311b50

Switch JAX to the new `ProgramShape::AddParameter()` API.

Google-ML-Automation committed 1 year ago

d3a84289

Merge pull request #28156 from cgarciae:fix-prngkey

Google-ML-Automation committed 1 year ago

c70b43ce

Merge pull request #28155 from Caslyn:fix-doc-link

Google-ML-Automation committed 1 year ago

59b2addc

Allow all reshapes if the operand is fully replicated

yashk2810 committed 1 year ago

36cd3137

[docs] fix typo in doc link

Caslyn committed 1 year ago

6a680941

Older