ROCm/jax

Pull Requests Commits

Merge pull request #24896 from hawkinsp:pad

Google-ML-Automation committed 1 year ago

fcde8aa6

Merge pull request #24751 from Stella-S-Yan:feature/default_device_str

Google-ML-Automation committed 1 year ago

cea81767

[sharding_in_types] Support shard_map with sharding in types. Right now only full manual mode is supported.

yashk2810 committed 1 year ago

05716b58

[Mosaic][TPU] Omit short circuiting of relayout (we should always relayout!) and implement product mismatch case for where we relayout from replicated to offset, and the number of vregs changes.

Google-ML-Automation committed 1 year ago

a8464ce7

Merge pull request #24865 from jakevdp:dep-clear-backends

Google-ML-Automation committed 1 year ago

89f411ab

Fix pgle test handling of directory removal.

dfm committed 1 year ago

de07be5c

Merge pull request #24899 from vfdev-5:patch-2

Google-ML-Automation committed 1 year ago

3380a211

finalize deprecation of jax.clear_backends

jakevdp committed 1 year ago

f401c979

Fixed typo in asan.yaml

vfdev-5 committed 1 year ago

Verified 6644bbc9

Merge pull request #23374 from jaro-sevcik:mock-topology-config

Google-ML-Automation committed 1 year ago

a1eb5cea

Make the jaxpr for jnp.pad in "constant" mode more succinct.

hawkinsp committed 1 year ago

ad5a0621

[pallas:mosaic_gpu] Ported two pipelining optimizations to `emit_pipeline`

superbobry committed 1 year ago

aefe6215

Merge pull request #23805 from dfm:ffi-examples-state

Google-ML-Automation committed 1 year ago

83700828

Merge pull request #24889 from skye:barrier

Google-ML-Automation committed 1 year ago

268c86f8

Merge pull request #24873 from jax-ml:no-gen-linear-util

Google-ML-Automation committed 1 year ago

d69fe1f8

Merge pull request #24810 from jakevdp:bitwise-reduce

Google-ML-Automation committed 1 year ago

842d93e1

Avoid repeatedly rebuilding a tuple in issubdtype.

hawkinsp committed 1 year ago

2463bf1f

Merge pull request #24886 from carlosgmartin:fix_typos

Google-ML-Automation committed 1 year ago

426e13a5

Merge pull request #24883 from jakevdp:doc-installation

Google-ML-Automation committed 1 year ago

a792fb06

Don't override --xla_tpu_use_enhanced_launch_barrier if explicitly set

skye committed 1 year ago

94f9a488

Stop using generators for linear_util transformations.

dougalm committed 1 year ago

1c9b23c5

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4a884d41

Merge pull request #24874 from pearu:pearu/square_p

Google-ML-Automation committed 1 year ago

14e08aa2

Fix typos: Change 'arugments' to 'arguments'.

carlosgmartin committed 1 year ago

307e88f2

doc: link directly to installation on the main page

jakevdp committed 1 year ago

72a4692b

Merge pull request #22899 from trevor-m:cache

Google-ML-Automation committed 1 year ago

0a755aec

Merge pull request #24069 from sergachev:cudnn_fusion_test_a100

Google-ML-Automation committed 1 year ago

12c8c68c

When caching is enabled, also enable XLA caching features as well

trevor-m committed 1 year ago

a79d307a

Add Pallas Triton lowering for jax.lax.bitcast_convert_type.

petebu committed 1 year ago

558ebb9f

pearu committed 1 year ago

4d0a007d

Older