ROCm/jax

Pull Requests Commits

Add missing CPython build deps for pyenv

mrodden committed 1 year ago

fafa03c6

Fix not finding wheels in bazel output

mrodden committed 1 year ago

701cda8e

Fix jenkins notty issue

mrodden committed 1 year ago

df2d140f

Add defaults for ROCm build vars

mrodden committed 1 year ago

319ebf81

Add copyright and license headers to new files

mrodden committed 1 year ago

abe44f6d

Add support for ROCm development builds

mrodden committed 1 year ago

a1a0a4ec

Add internal release support to get_rocm.py

mrodden committed 1 year ago

3175f13c

[ROCm] Change ROCm builds to manylinux wheels

mrodden committed 1 year ago

1e58d767

Register shutdown code at import to hopefully get registered before any other atexit callbacks.

BrianWieder committed 1 year ago

ee31e95e

Merge pull request #23014 from google:dependabot/github_actions/actions/upload-artifact-4.3.6

a-googler committed 1 year ago

7a873c03

Merge pull request #22979 from jakevdp:intersect1d-size

a-googler committed 1 year ago

da259f8d

Bump actions/upload-artifact from 4.3.5 to 4.3.6

dependabot[bot] committed 1 year ago

Verified 802abfef

Make custom partitioning work without a mesh context manager. If the arguments have NamedSharding on them, then inside `partition` function, we should get NamedSharding without the existence of the mesh context manager

yashk2810 committed 1 year ago

53045380

Add a `jax.process_indices` function.

dfm committed 1 year ago

60bf5b77

Support `None` leaves in arguments to gradient of a call_tf wrapped function.

brills committed 1 year ago

ad74e55d

Update shape polymorphism tests to skip lu_pivots_to_permutations tests when jaxlib version is too old.

dfm committed 1 year ago

4eb5ef28

Merge pull request #22997 from superbobry:maint-2

a-googler committed 1 year ago

112cae1d

[mosaic_gpu] Store untiled splat layout

cperivol committed 1 year ago

cd4e91b2

[pallas] Move ops_test.py from jax_triton to jax/pallas

gnecula committed 1 year ago

7f680aaa

Collapsed a few unnecessary ``if TYPE_CHECKING`` blocks

superbobry committed 1 year ago

c9142cbe

Add support for shape polymorphism with lu_pivots_to_permutation.

dfm committed 1 year ago

3c014a4c

Make `ffi_call` tests backwards compatible with the released jaxlib.

dfm committed 1 year ago

ae5b4284

Merge pull request #22667 from ROCm:rocm-jax-triton-add-get_arch_detail

a-googler committed 1 year ago

be4d52b8

[ROCm]: Add get_arch_details for triton kernel call

Rahul Batra committed 1 year ago

4b7c198a

Merge pull request #22991 from froystig:keys2

a-googler committed 1 year ago

d2fa8849

Merge pull request #22989 from froystig:keys

a-googler committed 1 year ago

312ebdb8

Merge pull request #22992 from froystig:docs2

a-googler committed 1 year ago

b9e80df6

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3c7bd54c

in `dot` docstring, format and link to `dot_general`

froystig committed 1 year ago

c54ffd41

emphasize typed over legacy RNG keys in `random` module docs

froystig committed 1 year ago

dd535d88

Older