ROCm/jax

Pull Requests Commits

Avoid querying metadata query to check if it's GCE if `TPU_SKIP_MDS_QUERY` is set.

Google-ML-Automation committed 1 year ago

84cd3567

[Pallas:TPU] Match lax.pow(float, int) behavior in Pallas.

WindQAQ committed 1 year ago

11faf680

Merge pull request #24482 from jakevdp:unwrap-doc

Google-ML-Automation committed 1 year ago

3c2c60e4

[Pallas] Allow core_map's mesh to discharge backend specific effects

apaszke committed 1 year ago

88d231a3

[Pallas:MGPU] Allow initializing accumulators with values in registers

apaszke committed 1 year ago

5b3b6e84

Add test for jax2tf conversion of dot general with algorithm.

dfm committed 1 year ago

5ea62154

[pallas:mosaic_gpu] An extremely specific heuristic to allow swiglu.

cperivol committed 1 year ago

40c92c1f

[pallas:mosaic_gpu] Memref reshape Transform to allow the user to reshape references.

cperivol committed 1 year ago

155aa6ca

Add an extension mechanism to run_state that allows:

sharadmv committed 1 year ago

ce8ecbd1

Better docs for jnp.unwrap

jakevdp committed 1 year ago

d6f4ce16

Merge pull request #24473 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

abc6c004

Merge branch 'release/0.4.35'

hawkinsp committed 1 year ago

2aeda178

Merge pull request #24446 from gnecula:export_doc

Google-ML-Automation committed 1 year ago

7ad73e44

Merge pull request #24469 from jakevdp:indices-doc

Google-ML-Automation committed 1 year ago

81bee421

Fix jax2tf failure coming from dot_general

yashk2810 committed 1 year ago

4688da31

Better documentation for jnp.indices

jakevdp committed 1 year ago

9038bb26

[sharding_in_types][Take 2] Add `out_type` argument to `einsum` and `dot_general` to allow specifying for the output type. Right now, it only accept a `NamedSharding` but in the future we can allow a polymorphic type of: `jax.ShapeDtypeStruct | Sharding | Layout`.

yashk2810 committed 1 year ago

f8a1f02d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

32be1992

JAX release 0.4.35

hawkinsp committed 1 year ago

81991d87

Use libtpu releases rather than libtpu-nightly for jax[tpu].

hawkinsp committed 1 year ago

e4f3f8f0

Merge pull request #24465 from jakevdp:fix-mypy

Google-ML-Automation committed 1 year ago

1c6b0a91

Merge pull request #21524 from andportnoy:aportnoy/unknown-platform-lowering-warning

Google-ML-Automation committed 1 year ago

9a2dd19a

Merge pull request #24452 from jakevdp:insert-doc

Google-ML-Automation committed 1 year ago

1e41d5ef

Merge pull request #24467 from andportnoy:patch-2

Google-ML-Automation committed 1 year ago

1a2737b7

jakevdp committed 1 year ago

84985021

Merge pull request #24463 from jakevdp:ml-dtypes-version

Google-ML-Automation committed 1 year ago

92aa9a39

[Pallas TPU] Fix lowering for `jnp.remainder`

ayaka14732 committed 1 year ago

c60bafcc

Add back the import of `jtu` in `flash_attention.py`

andportnoy committed 1 year ago

Verified 63789849

Raise an error when registering a lowering for an unknown platform

andportnoy committed 1 year ago

2aaa108f

[Pallas TPU] Add lowerings for bf16 `jnp.ceil` and `jnp.floor` in TPU v6+

ayaka14732 committed 1 year ago

2b7b0742

Newer Older