ROCm/jax

Pull Requests Commits

Merge pull request #27973 from mattjj:iadd-gotcha

Google-ML-Automation committed 1 year ago

1af747b3

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

mattjj committed 1 year ago

13c7183c

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

tlongeri committed 1 year ago

2e4c0ec7

Merge pull request #27868 from apaszke:mgpu-synchronization-docs

Google-ML-Automation committed 1 year ago

773b323b

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4fd610fc

Add reference docs for Pallas:MGPU synchronization primitives

apaszke committed 1 year ago

7edd5d50

Merge pull request #27981 from apaszke:mgpu-sm-image

Google-ML-Automation committed 1 year ago

a51307a5

Properly center and size the SM image in the GPU docs

apaszke committed 1 year ago

ca50cae5

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

69173a28

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

hawkinsp committed 1 year ago

c69e61e1

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

froystig committed 1 year ago

566d0775

revert making `Executable` an ABC

froystig committed 1 year ago

99ca1460

Remove various methods from `MeshExecutable`

yashk2810 committed 1 year ago

4ff78e6a

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

froystig committed 1 year ago

19d3d954

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

gnecula committed 1 year ago

dc102009

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

Google-ML-Automation committed 1 year ago

e1cad345

Rename is_closed to is_open in the shardy shardings

yashk2810 committed 1 year ago

8afc833c

Merge pull request #27970 from mattjj:while-readonly-carry-optimization

Google-ML-Automation committed 1 year ago

1a4a86aa

re-index jaxpr input effects in move_binders_to_front

mattjj committed 1 year ago

29f65f04

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

bythew3i committed 1 year ago

0fa732ea

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

mattjj committed 1 year ago

b3f49e42

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

Google-ML-Automation committed 1 year ago

b2a8df71

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

hawkinsp committed 1 year ago

6fc78a5a

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

pschuh committed 1 year ago

c0d97a68

Deprecate `PositionalSharding` and `GSPMDSharding`

yashk2810 committed 1 year ago

6efcf44b

Reverts 907725dfd7a7fb612c4f6d975bb462f1ae1a21d7

mattjj committed 1 year ago

e9364f4b

Rename TPU bazel test tags.

hawkinsp committed 1 year ago

904419cb

[Pallas] Allow 1D iota

justinjfu committed 1 year ago

27c07f7c

Use dash instead of underscore for extras.

nitins17 committed 1 year ago

5cf74cc7

document SPMD pipeline parallelism

Google-ML-Automation committed 1 year ago

8e9fca1d

Older