ROCm/jax

Pull Requests Commits

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

[mosaic_gpu] Fix signedness handling in FragmentedArray._pointwise.

petebu committed 1 year ago

1afb05e2

Merge pull request #24593 from froystig:random-dtypes

Google-ML-Automation committed 1 year ago

ae46b756

Add test_compute_on_host_shared_sharding in memories_test

Google-ML-Automation committed 1 year ago

4d60db17

represent `random.key_impl` of builtin RNGs by canonical string name

froystig committed 1 year ago

4bb81075

[Mosaic] Add `tpu.log` verification on SC

naummo committed 1 year ago

6c291d67

Fix a bug where constant deduplication used an inappropriate inequality.

hawkinsp committed 1 year ago

867a3618

Make deprecated jax.experimental.array_api module visibility internal-only

Jake VanderPlas committed 1 year ago

8c71d1ad

[Mosaic] Extend tpu.sem_signal with subcore_id

naummo committed 1 year ago

c04aec9d

Reverts 2075b091c4e83f0bdbd0d47812a72114fb8b937a

hawkinsp committed 1 year ago

525b646c

Move JAX example to public XLA:CPU API

changm committed 1 year ago

42fbd301

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3161a284

[Mosaic] Add target core type parameter to tpu.sem_signal

naummo committed 1 year ago

0d36b0b4

[pallas:mosaic_gpu] `copy_gmem_to_smem` no longer requires `barrier` to be a keyword argument

superbobry committed 1 year ago

1bf70fbb

Merge pull request #24970 from hawkinsp:split

Google-ML-Automation committed 1 year ago

2075b091

Add a new API jax.lax.split.

hawkinsp committed 1 year ago

2c80d1af

Add test utility for accessing jaxlib version tuple.

dfm committed 1 year ago

a59bbb7c

Add alternate implementation of threefry as a pallas kernel.

justinjfu committed 1 year ago

c44f11d1

[Mosaic TPU] Add general tpu.vector_store and support masked store.

bythew3i committed 1 year ago

6c31efa3

Newer Older