ROCm/jax

Pull Requests Commits

[ROCm] Only run core tests for ROCm post-build checks

charleshofer committed 362 days ago

d50206b8

Use an PartitionSpec.UNCONSTRAINED to represent unconstrained dimensions in ParsedPartitionSpec, rather than None.

hawkinsp committed 363 days ago

cf308a84

Update XLA dependency to use revision

Google-ML-Automation committed 364 days ago

8401d9b7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b3df4d58

Merge pull request #26407 from jakevdp:printoptions-doc

Google-ML-Automation committed 1 year ago

28903574

Merge pull request #26406 from jakevdp:lax-docs

Google-ML-Automation committed 1 year ago

a7c20e7f

Don't call `get_cur_mesh_sharding` if sharding-in-types mode is not enabled

yashk2810 committed 1 year ago

21e1be33

Speed up NamedSharding construction.

hawkinsp committed 1 year ago

f21b0f03

Increase the absolute error tolerance to fix flaky tests.

Rifur13 committed 1 year ago

cd075375

Merge pull request #26377 from mattjj:maintain-mutable-array-sharding

Google-ML-Automation committed 1 year ago

c0ba3626

Merge pull request #26376 from jakevdp:array-creation

Google-ML-Automation committed 1 year ago

ec477634

[pallas:mosaic_gpu] Fixed a bug in how `delay_release` is handled in `emit_pipeline`

superbobry committed 1 year ago

e5058079

[mutable-arrays] persist shardings through xla computations

mattjj committed 1 year ago

719031c1

Merge pull request #26409 from hawkinsp:fstring

Google-ML-Automation committed 1 year ago

7ffb613b

Merge pull request #26383 from jakevdp:jnp-sorting

Google-ML-Automation committed 1 year ago

3b470b95

DOC: make clear that printoptions are NumPy aliases

jakevdp committed 1 year ago

08563842

Fix a missing "f" on an f-string.

hawkinsp committed 1 year ago

d01520c6

Merge pull request #25810 from ROCm:gh-9948-add-gpu-ci-upstream

Google-ML-Automation committed 1 year ago

1ca8807d

jax.lax: better docs for hyperbolic trig functions

jakevdp committed 1 year ago

311e6683

[xla:python] Add a mechanism for "batch partitioning" of FFI calls.

dfm committed 1 year ago

c521bc62

Remove the unused cu_cholesky_update kernel in favor of the FFI version.

dfm committed 1 year ago

5bc17f7e

refactor: move sorting ops out of lax_numpy

jakevdp committed 1 year ago

d3b3cd36

Merge pull request #26344 from Cjkkkk:disable_head_256_on_bw

Google-ML-Automation committed 1 year ago

e56c7dc5

[pallas:triton] Really revert to the lowering using Triton IR

superbobry committed 1 year ago

35351f95

Merge pull request #26368 from ROCm:fix-rocm-pallas-lowerings

Google-ML-Automation committed 1 year ago

6ad38af4

[Mosaic GPU] Add tests for various tcgen05.mma configurations

apaszke committed 1 year ago

3b5e91b8

Add check to `lax.composite` to prevent `DynamicJaxprTracer` type errors.

ghpvnist committed 1 year ago

2828bce2

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4c402b7b

Fix file path regexp for Windows

gnecula committed 1 year ago

c1734f22

Update RNN kernels to use FFI.

dfm committed 1 year ago

c6e83903

Older