ROCm/jax

Pull Requests Commits

Merge pull request #25019 from jakevdp:lax-pad-doc

Google-ML-Automation committed 1 year ago

334bd4d0

DOC: add examples for jax.lax.pad

jakevdp committed 1 year ago

2699e950

[array api] use most recent version of array_api_tests

jakevdp committed 1 year ago

f749fca7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6fe78042

[sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

yashk2810 committed 1 year ago

9b941808

Merge pull request #25011 from jakevdp:jnp-linalg-module

Google-ML-Automation committed 1 year ago

d219439d

[Pallas TPU] Support masked store

bythew3i committed 1 year ago

9d2f62f8

[sharding_in_types] Make flash_attention forward pass in TPU pallas work nicely with sharding in types. Backward pass is still busted which I will fix in follow up CLs.

yashk2810 committed 1 year ago

40fc6598

Merge pull request #25013 from hawkinsp:relnotes

Google-ML-Automation committed 1 year ago

19b4996e

Mention stackless in the release notes.

hawkinsp committed 1 year ago

dfe27a16

Merge pull request #25008 from skye:barrier

Google-ML-Automation committed 1 year ago

1a3e693a

Set __module__ attribute of jax.numpy.linalg APIs

jakevdp committed 1 year ago

621e39de

[pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test

superbobry committed 1 year ago

9584ee3b

Don't psum over auto mesh dims in _unmentioned2.

pschuh committed 1 year ago

2c9b917b

Remove internal KeyArray alias

jakevdp committed 1 year ago

fee272e5

Merge pull request #25004 from jax-ml:linearize-trace

Google-ML-Automation committed 1 year ago

eab9026c

Make a direct linearize trace.

dougalm committed 1 year ago

d0f17c0c

[pallas mgpu] Lowering for while loops as long as they are secretly for loops.

cperivol committed 1 year ago

8d84f283

Fix KeyError recently introduced in cloud_tpu_init.py

skye committed 1 year ago

62225926

Merge pull request #25005 from jakevdp:py313

Google-ML-Automation committed 1 year ago

439d34da

Merge pull request #25007 from jakevdp:deps

Google-ML-Automation committed 1 year ago

800add2a

Deprecate several private APIs in jax.lib

jakevdp committed 1 year ago

85e2969a

Simplify handling of `DotAlgorithmPreset` output types.

chr1sj0nes committed 1 year ago

1e9e85a3

Mention python 3.13 in docs & package metadata

jakevdp committed 1 year ago

a4266b5e

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a582df02

[pallas] Do not skip vmap tests on GPU when x64 is enabled

superbobry committed 1 year ago

1df4b5f7

[mosaic_gpu] Handle older `jaxlib`s in the profiler module

superbobry committed 1 year ago

04e4c69f

[mosaic_gpu] Fixed `FragmentedArray` comparisons with literals

superbobry committed 1 year ago

f442d40f

[pallas:mosaic_gpu] `copy_smem_to_gmem` now supports `wait_read_only`

superbobry committed 1 year ago

c76e5fe9

[pallas:mosaic_gpu] Add Pallas Mosaic GPU lowering for jax.lax.bitcast_convert_type.

petebu committed 1 year ago

14da7ebb

Newer Older