ROCm/jax

Pull Requests Commits

[ROCm] build script fix

Ruturaj4 committed 1 year ago

a2824862

[Mosaic TPU] Support native int4 @ int4

bythew3i committed 1 year ago

a47b7556

Merge pull request #24305 from jakevdp:signbit-doc

Google-ML-Automation committed 1 year ago

94152e8d

Merge pull request #24294 from jakevdp:invert-doc

Google-ML-Automation committed 1 year ago

87d8f381

Fix typo in Pallas TPU matmul doc. I think the logical layout of the input array is non-transposed, rather than transposed?

18praveenb committed 1 year ago

3a3190fb

Merge pull request #23684 from simonster:sjk/fix-prefix-error

Google-ML-Automation committed 1 year ago

e461c049

Skip some Shardy-enabled tests if XLA < 292.

belitskiy committed 1 year ago

2f2fd8a3

Improve docs for jnp.invert and related functions

jakevdp committed 1 year ago

dd4a0408

Merge pull request #24251 from dfm:dot-algorithm-jax2tf

Google-ML-Automation committed 1 year ago

2c2c1eeb

[Pallas] Use core_map instead of shard_map for Shmallas

sharadmv committed 1 year ago

cd78c653

Add back missing manual axis to fallback.

pschuh committed 1 year ago

b0768906

Add a private property to NamedSharding called `_logical_device_ids` which allows you to pass a custom `tile_assignment_devices()` equivalent.

yashk2810 committed 1 year ago

2f6cb89a

Better docs for jnp.signbit

jakevdp committed 1 year ago

a0969868

Internal change.

WindQAQ committed 1 year ago

0788a0a5

Merge pull request #24299 from jakevdp:apply-doc

Google-ML-Automation committed 1 year ago

829e3159

Merge pull request #24258 from dfm:remove-jaxlib-version-checks

Google-ML-Automation committed 1 year ago

dfee4d15

Merge pull request #23467 from abhinavgoel95:patch-4

Google-ML-Automation committed 1 year ago

2a828d5d

[XLA:Collective] Expose a factory for constructing HLOSharding with explicit device ordering.

Tongfei-Guo committed 1 year ago

d621737f

[Pallas TPU] Add lowerings for scalar `sin`, `cos`, `tan` and `tanh`

ayaka14732 committed 1 year ago

bfc3d3cd

Better docs for jnp.apply_along_axis & apply_over_axes

jakevdp committed 1 year ago

0c307fe7

Add a memory saving index rewrite step to vmap with ragged inputs over pallas_call.

Google-ML-Automation committed 1 year ago

1f0b5728

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

fff3b874

Merge pull request #24290 from jax-ml:dependabot/github_actions/actions/upload-artifact-4.4.3

Google-ML-Automation committed 1 year ago

90cd8a79

Merge pull request #24289 from jax-ml:dependabot/github_actions/actions/cache-4.1.1

Google-ML-Automation committed 1 year ago

13bc4978

Bump actions/upload-artifact from 4.4.1 to 4.4.3

dependabot[bot] committed 1 year ago

Verified 93adc0e9

Bump actions/cache from 4.1.0 to 4.1.1

dependabot[bot] committed 1 year ago

Verified 0fdd6535

[Shardy] Inline meshes when using shardy and get rid of global meshes from the MLIR body.

yashk2810 committed 1 year ago

824ccd71

#sdy Run inlined mesh lifter pass at the end of JAX lowering.

bartchr808 committed 1 year ago

75e22f2c

Merge pull request #24271 from jakevdp:hist-doc

Google-ML-Automation committed 1 year ago

d15d70d6

Merge pull request #24264 from ROCm:ci_apt_update

Google-ML-Automation committed 1 year ago

1de9f25c

Older