ROCm/jax

Pull Requests Commits

Merge pull request #24912 from jakevdp:jnp-module

Google-ML-Automation committed 1 year ago

05d66d7c

Add new CI script for running Bazel GPU presubmits

nitins17 committed 1 year ago

14187399

Merge pull request #24932 from hawkinsp:gather

Google-ML-Automation committed 1 year ago

65f9c785

Add a GPU implementation of `lax.linalg.eig`.

dfm committed 1 year ago

ccb33170

Merge pull request #24933 from hawkinsp:pow

Google-ML-Automation committed 1 year ago

afdc7927

Merge pull request #24930 from hawkinsp:dedup

Google-ML-Automation committed 1 year ago

f7ae0f99

[AutoPGLE] Temporary disable pgle_test in the OSS.

Google-ML-Automation committed 1 year ago

ed250b89

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

742cabc5

Adds an env that can let users provide a custom version suffix for jax dev build.

yliu120 committed 1 year ago

27bf80a5

Use a direct StableHLO lowering for pow.

hawkinsp committed 1 year ago

8a6c560b

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7b9914d7

Return a ndarray in shape_as_value if the shape is known to be constant.

hawkinsp committed 1 year ago

1d519f4c

Deduplicate constants in StableHLO lowering.

hawkinsp committed 1 year ago

626aea01

Adds a flag to control proxy env checking.

yliu120 committed 1 year ago

609dfac2

[sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot

yashk2810 committed 1 year ago

8525ef2b

Merge pull request #24917 from emilyfertig:emilyaf-sharp-bits

Google-ML-Automation committed 1 year ago

efd23276

Consolidate material on PRNGs and add a short summary to Key Concepts.

emilyfertig committed 1 year ago

225a2a5f

Merge pull request #24914 from jakevdp:fix-pyi

Google-ML-Automation committed 1 year ago

1aa5de66

Merge pull request #24918 from emilyfertig:emilyaf-logical-op-example

Google-ML-Automation committed 1 year ago

605c6051

DOC: update main landing page style

barnesjoseph committed 1 year ago

81cdc882

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1780ff29

Add an example on logical operators to the tutorial.

emilyfertig committed 1 year ago

5f1e3f56

Add missing functions to jax.numpy type interface

jakevdp committed 1 year ago

5f942844

Merge pull request #24913 from hawkinsp:threefry

Google-ML-Automation committed 1 year ago

d8085008

Lower threefry as an out-of-line MLIR function on TPU.

hawkinsp committed 1 year ago

23e9142d

[Mosaic TPU] Support 1D concat: set implicit_dim to kSecondMinor to treat 1D (N,) as (1, N) and then tile it as (1, 128)

Google-ML-Automation committed 1 year ago

1471702a

Set __module__ attribute for objects in jax.numpy

jakevdp committed 1 year ago

f652b6ad

[sharding_in_types] Handle collective axes in lowering rules more generally. If any axis is collective, set all dims of aval to unspecified dims in `wrap_with_sharding_op`.

yashk2810 committed 1 year ago

9a0e9e55

Merge pull request #24862 from emilyfertig:emilyaf-control-flow-tutorial

Google-ML-Automation committed 1 year ago

4511f0c6

Merge pull request #24907 from jakevdp:array-api

Google-ML-Automation committed 1 year ago

1c31860f

Newer Older