ROCm/jax

Pull Requests Commits

Remove `DeviceLocalLayout` and `.device_local_layout` aliases.

yashk2810 committed 1 year ago

c057c7cd

Move jax._src.shard_map to its own BUILD rule

Jake VanderPlas committed 1 year ago

353363ae

Move jax._src.flatten_util to its own build rule

Jake VanderPlas committed 1 year ago

5749a34d

Finish `jax._src.lax` BUILD refactor

Jake VanderPlas committed 1 year ago

526ff073

Merge pull request #29626 from jakevdp:jax-numpy-imports

Google-ML-Automation committed 1 year ago

726a4094

Make canonicalization of reshapes more robust:

Google-ML-Automation committed 1 year ago

7d924e8f

Check for existence of all `None`s in a spec to consider a sharding as replicated in reshape.

yashk2810 committed 1 year ago

ab8720c2

Use `Layout` in place of `DeviceLocalLayout`

yashk2810 committed 1 year ago

88b0747a

Make `axis_types` argument to `Mesh`, `AbstractMesh` and `make_mesh` an optional positional argument. It can still be used as a keyword argument so it's backwards compatible.

yashk2810 committed 1 year ago

584f9bfb

Fix some sharding-in-types error messages

yashk2810 committed 1 year ago

53ff15a3

Go via device_put in reshard's impl rule if the mesh in context and the mesh is not multi-process

yashk2810 committed 1 year ago

2d0df4d7

Add pow to fusor.

Google-ML-Automation committed 1 year ago

4f0eeea6

Expose reshard, auto_axes and explicit_axes via the `jax.sharding` endpoint graduating them from `jax.experimental`.

yashk2810 committed 1 year ago

ff402188

Merge pull request #29565 from Cjkkkk:disable_cudnn_sdpa_on_gb300

Google-ML-Automation committed 1 year ago

48feb750

Merge pull request #29629 from mattjj:linear-solve-fix

Google-ML-Automation committed 1 year ago

1f75e16c

instantiate symbolic zero cotangents in _linear_solve_transpose_rule

mattjj committed 1 year ago

a26a1a72

jax._src.numpy: avoid top-level imports

jakevdp committed 1 year ago

5cafe62d

Expose local/global topology exchange timeouts for CPU client with collectives.

Google-ML-Automation committed 1 year ago

f141c365

[Pallas Fuser] Use partial_eval instead of DCE for fusible short-circuit check

jburnim committed 1 year ago

636691bb

[Pallas Fuser] Add support for custom_vjp functions

sharadmv committed 1 year ago

0320c354

Add REG type for GPUs

Google-ML-Automation committed 1 year ago

5f0b64e3

[Pallas][Mosaic GPU] Fix SMEM aliasing of RHS for tcgen05.mma

justinjfu committed 1 year ago

25c72d1f

[Mosaic:TPU] Extend support for ext to be symmetrical with trunc

tlongeri committed 1 year ago

2958740b

Merge pull request #29421 from mattjj:mutable-array-custom-vjp-scan2-again

Google-ML-Automation committed 1 year ago

64736c9b

Add more models to auto tune and update tuned block.

Google-ML-Automation committed 1 year ago

9f9a7d41

[Pallas:MGPU] Return the same grid as passed to MGPU from the Pallas:MGPU lowering

apaszke committed 1 year ago

234b87a8

[Take 2][Roll-forward] Remove Layout, .layout, .input_layouts and .output_layouts and replace it with Format, .format, .input_formats and .output_formats in JAX

yashk2810 committed 1 year ago

6e39a181

Reverts bfc07e2f10093ae169e2578872b04e8eb204a1d6

Google-ML-Automation committed 1 year ago

ea9fa64d

[pallas:mosaic] Aligned squeezed dimension inference for `tpu.memref_squeeze` with MLIR

superbobry committed 1 year ago

ddb783af

[Pallas:MGPU] Add a reference to the software pipelining guide

apaszke committed 1 year ago

22f7b7b5

Older