Code

CI: 01/08/25 upstream sync#196

Code

Merged

github-actions[bot] merged 34 commits intorocm-mainROCm/jax:rocm-mainfrom ci-upstream-sync-82_1ROCm/jax:ci-upstream-sync-82_1Copy head branch name to clipboardJan 8, 2025

+958-533

Commits

Commits on Dec 19, 2024

• Add activation offloading examples in a subsection

  zhenying-liu

  committed
  f124232

  Copy full SHA for f124232

• consolidate the code example

  zhenying-liu

  committed
  522a8fd

  Copy full SHA for 522a8fd

Commits on Jan 5, 2025

• Merge branch 'jax-ml:main' into activation-offloading-doc

  zhenying-liu

  authored
  5e3a692

  Copy full SHA for 5e3a692

Commits on Jan 7, 2025

• Update the advanced autodiff tutorial and replace some vmap with grad

  Exferro

  committed
  f881f50

  Copy full SHA for f881f50

• [Mosaic GPU] Use num_q_heads=2 in flash_attention.py

  Show description for 5b80892
  andportnoy

  authored
  5b80892

  Copy full SHA for 5b80892

• Clean up backend_or_name vs. platforms in lowering code.

  Show description for fdb6af8
  

  gnecula

  authored and

  Google-ML-Automation

  committed
  fdb6af8

  Copy full SHA for fdb6af8

• Merge pull request jax-ml#25441 from Exferro:fixed_advanced_autodiff_doc

  Show description for 57c2afe
  Google-ML-Automation

  committed
  57c2afe

  Copy full SHA for 57c2afe

• [AutoPGLE] FIx PGLE kokoro test failures.

  Show description for 4023810
  Google-ML-Automation

  committed
  4023810

  Copy full SHA for 4023810

• [Pallas] Fix pallas_call lowering mutating compiler params during Triton lowering.

  Show description for 8c9a539
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  8c9a539

  Copy full SHA for 8c9a539

• Update XLA dependency to use revision

  Show description for 00c363e
  Google-ML-Automation

  committed
  00c363e

  Copy full SHA for 00c363e

• Add an example demonstrating input-output aliasing with the FFI.

  dfm

  committed
  62656b3

  Copy full SHA for 62656b3

• Sort manual axes when lowering jax.shard_map to sdy.manual_computation, which ensures the determinism in the generated sdy.manual_computation.

  Show description for 64c0f62
  

  ZixuanJiang

  authored and

  Google-ML-Automation

  committed
  64c0f62

  Copy full SHA for 64c0f62

• Merge pull request jax-ml#25042 from dfm:ffi-example-input-output-alias

  Show description for f1777d5
  Google-ML-Automation

  committed
  f1777d5

  Copy full SHA for f1777d5

• [array api] update test suite to latest commit

  jakevdp

  committed
  f6c9e87

  Copy full SHA for f6c9e87

Commits on Jan 8, 2025

• Increase the minimum SciPy version to 1.11.1.

  Show description for 392a851
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  392a851

  Copy full SHA for 392a851

• [Pallas] Improvements to core_map

  Show description for 7be127f
  

  sharadmv

  authored and

  Google-ML-Automation

  committed
  7be127f

  Copy full SHA for 7be127f

• [sharding_in_types] Aval sharding under full auto mode should contain None and not UNCONSTRAINED because axis_types + pspec give the full picture.

  Show description for 755d6cd
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  755d6cd

  Copy full SHA for 755d6cd

• Merge pull request jax-ml#25761 from jakevdp:array-api-update

  Show description for 6d08f36
  Google-ML-Automation

  committed
  6d08f36

  Copy full SHA for 6d08f36

• Merge branch 'jax-ml:main' into activation-offloading-doc

  zhenying-liu

  authored
  21fb171

  Copy full SHA for 21fb171

• Add JAX events that have time spans, not only durations.

  Show description for 1bd781d
  Google-ML-Automation

  committed
  1bd781d

  Copy full SHA for 1bd781d

• Merge pull request jax-ml#25594 from zhenying-liu:activation-offloading-doc

  Show description for 81db321
  Google-ML-Automation

  committed
  81db321

  Copy full SHA for 81db321

• Removed leftover mentions of xmap from the code

  Show description for 90201ce
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  90201ce

  Copy full SHA for 90201ce

• Merge pull request jax-ml#25754 from andportnoy:patch-4

  Show description for 4718121
  Google-ML-Automation

  committed
  4718121

  Copy full SHA for 4718121

• [Mosaic] Use tpu::CreateMask for getX32VmaskByPaddingEnd.

  Show description for bf94389
  

  WindQAQ

  authored and

  Google-ML-Automation

  committed
  bf94389

  Copy full SHA for bf94389

• [Mosaic TPU] Add support for true divide in bf16 on TPUv6

  Show description for e954930
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  e954930

  Copy full SHA for e954930

• [Mosaic TPU] Add support for second minor broadcasts with packed types

  Show description for 5fd1b2f
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  5fd1b2f

  Copy full SHA for 5fd1b2f

• [Mosaic TPU] Be much more aggressive in inferring large 2nd minor layouts for 16-bit types on v6

  Show description for f96339b
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  f96339b

  Copy full SHA for f96339b

• [JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.

  Show description for 51b9fe3
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  51b9fe3

  Copy full SHA for 51b9fe3

• [pallas:mosaic_gpu] Fix the tests following the changes to pl.core_map

  Show description for f1f98af
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  f1f98af

  Copy full SHA for f1f98af

• Port tests away from setUpClass and setUpModule to setUp alone.

  Show description for 3fa5572
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  3fa5572

  Copy full SHA for 3fa5572

• Add a unittest test extension that runs test cases in parallel using threads.

  Show description for 0389d61
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  0389d61

  Copy full SHA for 0389d61

• #sdy Move Shardy mesh lift inlining pass after verification.

  Show description for 5c097c8
  

  bartchr808

  authored and

  Google-ML-Automation

  committed
  5c097c8

  Copy full SHA for 5c097c8

• Update XLA dependency to use revision

  Show description for 5511949
  Google-ML-Automation

  committed
  5511949

  Copy full SHA for 5511949

• Use our own XLA

  charleshofer

  committed
  8d94998

  Copy full SHA for 8d94998