ROCm/jax

Pull Requests Commits

Merge pull request #25489 from emilyfertig:distributed-init-input-validation

Google-ML-Automation committed 1 year ago

64d3871e

Delete unused internal symbols

Jake VanderPlas committed 1 year ago

482a6e73

Add more input validation to jax.distributed.initialize.

emilyfertig committed 1 year ago

a24e7032

[Mosaic] Extend macros to handle parentheses.

WindQAQ committed 1 year ago

60ebde89

[Mosaic] Remove TODOs that are already addressed or obsolete.

WindQAQ committed 1 year ago

77f3c114

Use common set of build options when building jaxlib+plugin artifacts together

nitins17 committed 1 year ago

6b096b0c

Ensure that the two offsets of a dynamic_slice have the same dtype regardless

bixia1 committed 1 year ago

16712b51

[Mosaic TPU] Add support for sqrt and rsqrt in bf16 on TPUv6

apaszke committed 1 year ago

de8fa8fd

[Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths

tlongeri committed 1 year ago

307c8d3a

[Mosaic GPU] Prototype of a warp-specialized pipeline emitter for Mosaic GPU.

justinjfu committed 1 year ago

d1294385

add test for partial-auto ppermute

mattjj committed 1 year ago

9f42b99a

Migrate shardy dialect extension to nanobind.

hawkinsp committed 1 year ago

9c3365fb

Merge pull request #25614 from jakevdp:dep-shaped-abstractify

Google-ML-Automation committed 1 year ago

3e2f2aab

DOC: clarify API compatibility discussion & mention EffVer

jakevdp committed 1 year ago

26930366

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1b385f38

Merge pull request #25606 from dfm:fft-nd

Google-ML-Automation committed 1 year ago

dc8b786d

Merge pull request #25440 from ROCm:jax_pypi-upstream

Google-ML-Automation committed 1 year ago

bad6f0f0

Merge pull request #25148 from ROCm:ci_install_sys_libs-upstream

Google-ML-Automation committed 1 year ago

25713f95

Add experimental support for building JAX CPU and GPU wheels with GCC.

Google-ML-Automation committed 1 year ago

c78ca042

Merge pull request #25516 from jakevdp:jep-effver

Google-ML-Automation committed 1 year ago

49c92460

Merge pull request #25590 from jakevdp:fix-one-hot-float

Google-ML-Automation committed 1 year ago

8402a988

Update documentation and add setup.py pypi bindings

Ruturaj4 committed 1 year ago

fefebebe

Always suppress the differing_executors Hypothesis health check

apaszke committed 1 year ago

23000a38

[Mosaic GPU] Commit to using `Vector`s everywhere (and no `Tensor`s).

bchetioui committed 1 year ago

3915f4a1

[jaxlib] Added a missing pytype_dep to :_triton_ext

superbobry committed 1 year ago

c4fae4a7

Add support for N-D FFTs with D>3.

dfm committed 1 year ago

c6131ee5

[pallas:triton] Fixed a typo in a type annotation

superbobry committed 1 year ago

af7a31f1

jax.nn.one_hot: deprecate non-integer inputs

jakevdp committed 1 year ago

8c3c441e

Remove internal uses of api_util.shaped_abstractify

jakevdp committed 1 year ago

5dc37d3f

[ROCm] ci build and dockerfile changes

Ruturaj4 committed 1 year ago

bfcace49

Newer Older