ROCm/jax

Pull Requests Commits

jnp.einsum: mention default value for optimize param

jakevdp committed 2 years ago

fbacebc1

CI: update pre-commit hooks to latest version

jakevdp committed 2 years ago

389eb97a

Update XLA dependency to use revision

a-googler committed 2 years ago

b093c4c9

Merge pull request #18274 from superbobry:no-config-methods

a-googler committed 2 years ago

21571fbb

Deprecated `define_*` and `DEFINE_*` methods on `jax.config`

superbobry committed 2 years ago

fd3a8b2c

Update XLA dependency to use revision

a-googler committed 2 years ago

1959d662

Update XLA dependency to use revision

a-googler committed 2 years ago

33a6863e

Do not init backends from topology construction, instead directly init the

pschuh committed 2 years ago

19c65353

Places the `remat` decorator on top of the body function.

Yunlong Liu committed 2 years ago

b99958db

Don't print all the devices in the mesh during ResourceEnv's repr. Just print the mesh shape.

yashk2810 committed 2 years ago

8ee58117

Merge pull request #18305 from superbobry:deprecate-jax-config-submodule

a-googler committed 2 years ago

69d0a404

Bump LIBTPU_OLDEST_VERSION_DATE to 20230915 to include https://github.com/openxla/xla/pull/5617.

Jieying Luo committed 2 years ago

8ffc88c1

Invalidate in-memory caches on XLA-AutoFDO profile version change.

a-googler committed 2 years ago

9ba305cc

Deprecated accessing `config` via the `jax.config` submodule

superbobry committed 2 years ago

c90f1f0c

[PJRT:C] Implement PjRtCApiClient::GetTopologyDescription

skye committed 2 years ago

58c86064

[JAX] Add an option `subset_by_index` that allows computing a contiguous subset of eigenvalues from eigh.

a-googler committed 2 years ago

11c4e2c8

Add mitigation techniques in the error message when a barrier timeout occurs.

yashk2810 committed 2 years ago

0d57330f

[Mosaic] Allow arbitrary broadcasts on non-tiled dimensions

apaszke committed 2 years ago

9752ba97

Update XLA dependency to use revision

a-googler committed 2 years ago

f438dac4

Fix an error format issue in jax.random

a-googler committed 2 years ago

15e6d7a3

Add nightly install instructions to docs

yashk2810 committed 2 years ago

87c05610

Use TSL's import of DUCC.

cantonios committed 2 years ago

873cffc7

[Mosaic] Add support for (almost) arbitrary reductions

apaszke committed 2 years ago

cd177fd5

[Pallas] Add Slice to public interface

a-googler committed 2 years ago

38e28693

Update XLA dependency to use revision

a-googler committed 2 years ago

f872ad78

AOT sharding mismatch error shouldn't have GSPMDSharding in it.

yashk2810 committed 2 years ago

4d153755

Remove bitwise conversions from _canonicalize_float_for_sort

majnemer committed 2 years ago

ba9fd774

Cleanup the handling of single- and multi-platform lowering in ModuleContext

gnecula committed 2 years ago

edbe49fb

Merge pull request #18277 from jakevdp:fix-rtd

a-googler committed 2 years ago

468f6667

DOC: set navigation_with_keys=False

jakevdp committed 2 years ago

a2688638

Newer Older