ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

73a015c4

Remove `eval_shape` as a method on Traced and rather use `.out_info` because `.trace` already evals.

yashk2810 committed 1 year ago

404981fd

Support _get_physical_tpu_mesh for TPU7X.

pschuh committed 1 year ago

dfa65845

[pallas:triton] Migrated example kernels to `plgpu.{load,store}`

superbobry committed 1 year ago

9c76aa51

[Mosaic][SC] Add custom assembly format to tpu.enqueue_indirect_dma

naummo committed 1 year ago

17e5b94f

[Mosaic GPU] Explicitly set kernel_name on both Pallas and plain Mosaic GPU codepaths.

justinjfu committed 1 year ago

07458292

[jax:custom_partitioning] Allow factors for non-batching dimensions to

bixia1 committed 1 year ago

5ef7deb7

Merge pull request #30258 from andportnoy:aportnoy/skip-tests-too-large-for-rtx-6000-pro

Google-ML-Automation committed 1 year ago

e7a3298d

[JAX] Move 7x mesh creation code to OSS.

yashk2810 committed 1 year ago

e62b4ab1

Make `tiling` argument of `Layout` class public and the property too. So basically it's `s/_tiling/tiling/g`.

yashk2810 committed 1 year ago

cbecf854

`lax.zeros_like_array` is no longer public

superbobry committed 1 year ago

37977251

Remove `local_config_nvshmem` repository from JAX WORKSPACE file.

Google-ML-Automation committed 1 year ago

c634a5e8

Replace the docker image with pre-installed CUDA with a clean Docker image.

Google-ML-Automation committed 1 year ago

f538df4d

[Mosaic] Allow vector::Extract for non-32 bits vector result.

WindQAQ committed 1 year ago

b09724a9

[bazel] Create jax/_src/BUILD and move _src build declarations

Jake VanderPlas committed 1 year ago

bd9ad438

Remove `--@local_config_nvshmem//:include_nvshmem_libs` usage and update mosaic GPU code to conform to the new `_solib` folder name.

Google-ML-Automation committed 1 year ago

6c0d816e

[pallas] Forked `load` and `store` into `triton` and `tpu`

superbobry committed 1 year ago

907a64d8

Merge pull request #30456 from daskol:fix/jaxlib/build-issue

Google-ML-Automation committed 1 year ago

a24ae9e9

[pallas] Moved atomic APIs to `jax.experimental.pallas.triton`

superbobry committed 1 year ago

77a1ddd7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

943126ca

[build] avoid top-level dependencies in jax/_src

Jake VanderPlas committed 1 year ago

e485ac95

Refactor: Add carry support to nd_loop

cperivol committed 1 year ago

718dba9b

Allow fusing all stride == 1 to support jax.lax.slice_in_dim in @fuser.fuse

Google-ML-Automation committed 1 year ago

a77cbc83

before iterating, copy dict that can be popped from via ARC / object deletion

mattjj committed 1 year ago

0d3feb8e

Merge pull request #30464 from mattjj:psend-precv-docs

Google-ML-Automation committed 1 year ago

ba3bfcae

add psend/precv to auto-generated docs

mattjj committed 1 year ago

54c618c9

[Mosaic GPU] Add missing namespace specified

daskol committed 1 year ago

Verified 133300f6

Merge pull request #30447 from justinjfu:mgpu_docs_fixes

Google-ML-Automation committed 1 year ago

b32721d4

[Mosaic GPU] Fix some typos in docs

justinjfu committed 1 year ago

98665a7f

Merge pull request #30420 from melissawm:refactor-httij

Google-ML-Automation committed 1 year ago

535829d7

Older