Commits

Branch selector

 main

User selector

All users

Datepicker

All time

Commit History

Commits on Jul 25, 2025

• Update XLA dependency to use revision

  Show description for 73a015c
  Google-ML-Automation

  committed
  73a015c

  Copy full SHA for 73a015c

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

  Show description for 404981f
  

  2 peopleauthored and

  Google-ML-Automation

  committed
  404981f

  Copy full SHA for 404981f

Commits on Jul 24, 2025

• Support _get_physical_tpu_mesh for TPU7X.

  Show description for dfa6584
  

  pschuh

  authored and

  Google-ML-Automation

  committed
  dfa6584

  Copy full SHA for dfa6584

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

  Show description for 9c76aa5
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  9c76aa5

  Copy full SHA for 9c76aa5

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

  Show description for 17e5b94
  

  naummo

  authored and

  Google-ML-Automation

  committed
  17e5b94

  Copy full SHA for 17e5b94

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

  Show description for 0745829
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  0745829

  Copy full SHA for 0745829

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

  Show description for 5ef7deb
  

  bixia1

  authored and

  Google-ML-Automation

  committed
  5ef7deb

  Copy full SHA for 5ef7deb

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

  Show description for e7a3298
  Google-ML-Automation

  committed
  e7a3298

  Copy full SHA for e7a3298

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

  Show description for e62b4ab
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  e62b4ab

  Copy full SHA for e62b4ab

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

  Show description for cbecf85
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  cbecf85

  Copy full SHA for cbecf85

• `lax.zeros_like_array` is no longer public

  Show description for 3797725
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  3797725

  Copy full SHA for 3797725

• Remove `local_config_nvshmem` repository from JAX WORKSPACE file.

  Show description for c634a5e
  Google-ML-Automation

  committed
  c634a5e

  Copy full SHA for c634a5e

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

  Show description for f538df4
  Google-ML-Automation

  committed
  f538df4

  Copy full SHA for f538df4

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

  Show description for b09724a
  

  WindQAQ

  authored and

  Google-ML-Automation

  committed
  b09724a

  Copy full SHA for b09724a

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

  Show description for bd9ad43
  

  Jake VanderPlas

  authored and

  Google-ML-Automation

  committed
  bd9ad43

  Copy full SHA for bd9ad43

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

  Show description for 6c0d816
  Google-ML-Automation

  committed
  6c0d816

  Copy full SHA for 6c0d816

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

  Show description for 907a64d
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  907a64d

  Copy full SHA for 907a64d

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

  Show description for a24ae9e
  Google-ML-Automation

  committed
  a24ae9e

  Copy full SHA for a24ae9e

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

  Show description for 77a1ddd
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  77a1ddd

  Copy full SHA for 77a1ddd

• Update XLA dependency to use revision

  Show description for 943126c
  Google-ML-Automation

  committed
  943126c

  Copy full SHA for 943126c

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

  Show description for e485ac9
  

  Jake VanderPlas

  authored and

  Google-ML-Automation

  committed
  e485ac9

  Copy full SHA for e485ac9

• Refactor: Add carry support to nd_loop

  Show description for 718dba9
  

  cperivol

  authored and

  Google-ML-Automation

  committed
  718dba9

  Copy full SHA for 718dba9

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

  Show description for a77cbc8
  Google-ML-Automation

  committed
  a77cbc8

  Copy full SHA for a77cbc8

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

  Show description for 0d3feb8
  

  mattjj

  authored and

  Google-ML-Automation

  committed
  0d3feb8

  Copy full SHA for 0d3feb8

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

  Show description for ba3bfca
  Google-ML-Automation

  committed
  ba3bfca

  Copy full SHA for ba3bfca

• add psend/precv to auto-generated docs

  mattjj

  committed
  54c618c

  Copy full SHA for 54c618c

Commits on Jul 23, 2025

• [Mosaic GPU] Add missing namespace specified

  daskol

  committed
  133300f

  Copy full SHA for 133300f

• Merge pull request #30447 from justinjfu:mgpu_docs_fixes

  Show description for b32721d
  Google-ML-Automation

  committed
  b32721d

  Copy full SHA for b32721d

• [Mosaic GPU] Fix some typos in docs

  justinjfu

  committed
  98665a7

  Copy full SHA for 98665a7

• Merge pull request #30420 from melissawm:refactor-httij

  Show description for 535829d
  Google-ML-Automation

  committed
  535829d

  Copy full SHA for 535829d

• Merge pull request #30446 from jakevdp:source-info-util-build

  Show description for 9b47a90
  Google-ML-Automation

  committed
  9b47a90

  Copy full SHA for 9b47a90

• Cache `canonicalize_sharding` and remove the canonicalization in `_scatter_update` since the out_sharding passed to scatter_update should already be

  Show description for 768a64e
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  768a64e

  Copy full SHA for 768a64e

• avoid top-level imports in source_info_util

  jakevdp

  committed
  338b4eb

  Copy full SHA for 338b4eb

• Skip Pallas and Mosaic GPU tests that don't fit on RTX 6000 PRO

  andportnoy

  committed
  240c2d0

  Copy full SHA for 240c2d0

• Test B200 against CUDA 12.8 only

  Show description for fd8a79f
  

  nitins17

  authored and

  Google-ML-Automation

  committed
  fd8a79f

  Copy full SHA for fd8a79f

Pagination

Previous

Next