ROCm/jax

Pull Requests Commits

Fix typo in jax.lax.linalg.symmetric_product description

ywrt committed 1 year ago

3d69db46

Remove accidental tab characters from Pallas:MGPU docs

apaszke committed 1 year ago

8a7164fe

Fix grep for label on Read the Docs.

dfm committed 1 year ago

d2dcdde1

Add Pallas:MGPU documentation for WGMMA

apaszke committed 1 year ago

55f96aaa

[explain-cache-miss] Improve tracing-cache-miss explanations

gnecula committed 1 year ago

0ebf2c67

add a brief description of the jax.Array-has-no-__iadd__ gotcha

mattjj committed 1 year ago

51d78ff9

[Mosaic:TPU] Add some invariant checking in VectorLayout ctor

tlongeri committed 1 year ago

1b7cbcb2

Add reference docs for Pallas:MGPU synchronization primitives

apaszke committed 1 year ago

157b36c7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9bf7b4cb

Properly center and size the SM image in the GPU docs

apaszke committed 1 year ago

b59c7fbd

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

72838996

Remove jax.lib.xla_client.{XlaComputation,Shape}.

hawkinsp committed 1 year ago

8ef47dfb

unify `stages.Lowering` and `stages.XlaLowering`

froystig committed 1 year ago

aeb4d190

revert making `Executable` an ABC

froystig committed 1 year ago

560de2eb

Remove various methods from `MeshExecutable`

yashk2810 committed 1 year ago

7af64122

unify `stages.Executable` and `stages.XlaExecutable`

froystig committed 1 year ago

cc89063f

[explain-cache-miss] Improve the detection of user file names

gnecula committed 1 year ago

592979c2

Add `ChunkedCausalMask` for Splash Attention to support attention masking similar to Llama4.

Google-ML-Automation committed 1 year ago

39656355

Rename is_closed to is_open in the shardy shardings

yashk2810 committed 1 year ago

d1238844

re-index jaxpr input effects in move_binders_to_front

mattjj committed 1 year ago

d20c6ae9

Re-landing #27937 with fewer bugs and more tests.

mattjj committed 1 year ago

df5d0ba8

[ragged-paged-attn][NFC] Make validate_inputs functions take same inputs as attention call.

bythew3i committed 1 year ago

552f62cc

Add the `method` argument to `jax.numpy.isin` stub.

Google-ML-Automation committed 1 year ago

724c7127

Deprecate jax.lax.infeed and jax.lax.outfeed.

hawkinsp committed 1 year ago

e459d3bb

Removed type annotations appear to be used and actually defined in python as a patch, rolling back.

pschuh committed 1 year ago

100eb7c0

Deprecate `PositionalSharding` and `GSPMDSharding`

yashk2810 committed 1 year ago

e43a7ac0

Reverts 907725dfd7a7fb612c4f6d975bb462f1ae1a21d7

mattjj committed 1 year ago

4b01a468

Rename TPU bazel test tags.

hawkinsp committed 1 year ago

8af10035

[Pallas] Allow 1D iota

justinjfu committed 1 year ago

527d9cb9

Use dash instead of underscore for extras.

nitins17 committed 1 year ago

d45ae60f

Newer Older