ROCm/jax

Pull Requests Commits

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

charleshofer committed 342 days ago

d50206b8

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

hawkinsp committed 343 days ago

cf308a84

Update XLA dependency to use revision

Google-ML-Automation committed 343 days ago

8401d9b7

Update XLA dependency to use revision

Google-ML-Automation committed 344 days ago

b3df4d58

Merge pull request #26407 from jakevdp:printoptions-doc

Google-ML-Automation committed 345 days ago

28903574

Merge pull request #26406 from jakevdp:lax-docs

Google-ML-Automation committed 345 days ago

a7c20e7f

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

yashk2810 committed 345 days ago

21e1be33

Speed up NamedSharding construction.

hawkinsp committed 345 days ago

f21b0f03

Increase the absolute error tolerance to fix flaky tests.

Rifur13 committed 345 days ago

cd075375

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

Google-ML-Automation committed 345 days ago

c0ba3626

Merge pull request #26376 from jakevdp:array-creation

Google-ML-Automation committed 345 days ago

ec477634

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

superbobry committed 345 days ago

e5058079

[mutable-arrays] persist shardings through xla computations

mattjj committed 345 days ago

719031c1

Merge pull request #26409 from hawkinsp:fstring

Google-ML-Automation committed 345 days ago

7ffb613b

Merge pull request #26383 from jakevdp:jnp-sorting

Google-ML-Automation committed 345 days ago

3b470b95

DOC: make clear that printoptions are NumPy aliases

jakevdp committed 345 days ago

08563842

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

hawkinsp committed 345 days ago

d01520c6

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

Google-ML-Automation committed 345 days ago

1ca8807d

jax.lax: better docs for hyperbolic trig functions

jakevdp committed 345 days ago

311e6683

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

dfm committed 345 days ago

c521bc62

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

dfm committed 345 days ago

5bc17f7e

refactor: move sorting ops out of lax_numpy

jakevdp committed 345 days ago

d3b3cd36

Merge pull request #26344 from Cjkkkk:disable_head_256_on_bw

Google-ML-Automation committed 345 days ago

e56c7dc5

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

superbobry committed 345 days ago

35351f95

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

Google-ML-Automation committed 345 days ago

6ad38af4

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

apaszke committed 345 days ago

3b5e91b8

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

ghpvnist committed 345 days ago

2828bce2

Update XLA dependency to use revision

Google-ML-Automation committed 345 days ago

4c402b7b

Fix file path regexp for Windows

gnecula committed 345 days ago

c1734f22

Update RNN kernels to use FFI.

dfm committed 346 days ago

c6e83903

Older