Code

CI: 01/14/25 upstream sync#202

Code

Open

github-actions[bot] wants to merge 74 commits intorocm-mainROCm/jax:rocm-mainfrom ci-upstream-sync-87_1ROCm/jax:ci-upstream-sync-87_1Copy head branch name to clipboard

+4,801-1,488

Commits

Commits on Jan 8, 2025

• fix grad(logsumexp) to produce 0s where where is False

  mattjj

  committed
  f0392a1

  Copy full SHA for f0392a1

Commits on Jan 9, 2025

• Move _reduce_window docstring to public func lax.reduce_window.

  cherrywoods

  authored
  6e9a34f

  Copy full SHA for 6e9a34f

• Fix struct string encoding non-determinism in the RNN descriptor.

  Show description for f0e1c3c
  sergachev

  committed
  f0e1c3c

  Copy full SHA for f0e1c3c

• [NFC] Refactor conversion lowering for Mosaic TPU

  Show description for f23979e
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  f23979e

  Copy full SHA for f23979e

• [Mosaic TPU] Fix a bug in the impl of sublane broadcasts for int8 and int4

  Show description for 07f4fd3
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  07f4fd3

  Copy full SHA for 07f4fd3

• Add a discussion of sharding to the FFI tutorial.

  dfm

  committed
  7294180

  Copy full SHA for 7294180

• Switch to a new thread-safe utility for catching warnings.

  Show description for b06779b
  hawkinsp

  committed
  b06779b

  Copy full SHA for b06779b

• Update XLA dependency to use revision

  Show description for e571694
  Google-ML-Automation

  committed
  e571694

  Copy full SHA for e571694

• Merge pull request jax-ml#25801 from cherrywoods:reduce_window_docstring

  Show description for 9cbd70b
  Google-ML-Automation

  committed
  9cbd70b

  Copy full SHA for 9cbd70b

• Merge pull request jax-ml#25626 from hawkinsp:warnings2

  Show description for 7718ac5
  Google-ML-Automation

  committed
  7718ac5

  Copy full SHA for 7718ac5

• Add sph_harm_y to jax.scipy.special and deprecate sph_harm.

  dfm

  committed
  5f3e0d9

  Copy full SHA for 5f3e0d9

• Merge pull request jax-ml#25753 from dfm:shp-harm-y

  Show description for 6c8b02d
  Google-ML-Automation

  committed
  6c8b02d

  Copy full SHA for 6c8b02d

• Clarify documentation of composites.

  Show description for 93ef0f1
  

  ghpvnist

  authored and

  Google-ML-Automation

  committed
  93ef0f1

  Copy full SHA for 93ef0f1

• Enable pivoted QR on CPU devices.

  Show description for c89be05
  tttc3

  committed
  c89be05

  Copy full SHA for c89be05

• #sdy enable pure callbacks and debug prints in JAX.

  Show description for dc53c56
  

  bartchr808

  authored and

  Google-ML-Automation

  committed
  dc53c56

  Copy full SHA for dc53c56

• Merge pull request jax-ml#25803 from sergachev:fix_rnn_desc

  Show description for 061408a
  Google-ML-Automation

  committed
  061408a

  Copy full SHA for 061408a

Commits on Jan 10, 2025

• Remove extraneous print statement in a test

  Show description for 6319126
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6319126

  Copy full SHA for 6319126

• [aot] Add support for as_text(debug_info=True).

  Show description for dd0447a
  gnecula

  committed
  dd0447a

  Copy full SHA for dd0447a

• [better_errors] Improve error message for lax.switch branches output structure mismatch

  Show description for c2adfbf
  gnecula

  committed
  c2adfbf

  Copy full SHA for c2adfbf

• Merge pull request jax-ml#25797 from gnecula:print_debug_info

  Show description for 228d3ce
  Google-ML-Automation

  committed
  228d3ce

  Copy full SHA for 228d3ce

• Merge pull request jax-ml#25800 from gnecula:improve_error_switch

  Show description for 8c23689
  Google-ML-Automation

  committed
  8c23689

  Copy full SHA for 8c23689

• Reverts dbe9ccd

  Show description for a27566c
  

  chr1sj0nes

  authored and

  Google-ML-Automation

  committed
  a27566c

  Copy full SHA for a27566c

• [Pallas] Improve testing for lowering of dtype conversions + fix uncovered bugs

  Show description for 74cf67d
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  74cf67d

  Copy full SHA for 74cf67d

• [pallas:mosaic_gpu] Tests now pass with x64 enabled

  Show description for 18018d9
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  18018d9

  Copy full SHA for 18018d9

• Skip RnnTest.test_struct_encoding_determinism.

  Show description for 86643a1
  

  belitskiy

  authored and

  Google-ML-Automation

  committed
  86643a1

  Copy full SHA for 86643a1

• Make JAX test suite pass (at least most of the time) with multiple threads enabled.

  Show description for c61b2f6
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  c61b2f6

  Copy full SHA for c61b2f6

• [Mosaic TPU] Add support for integer truncation from packed types

  Show description for d2a5e8d
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  d2a5e8d

  Copy full SHA for d2a5e8d

• Merge pull request jax-ml#25771 from dfm:custom-partition-ffi-tutorial

  Show description for 1fe72ee
  Google-ML-Automation

  committed
  1fe72ee

  Copy full SHA for 1fe72ee

• Avoid creating float0 JVPTracers

  dougalm

  committed
  ba9b2ca

  Copy full SHA for ba9b2ca

• Merge pull request jax-ml#20282 from tttc3:pivoted-qr

  Show description for 564b6b0
  Google-ML-Automation

  committed
  564b6b0

  Copy full SHA for 564b6b0

• [Mosaic GPU] Add initial layout mismatch resolution for splat/strided layouts.

  Show description for 1893881
  

  bchetioui

  authored and

  Google-ML-Automation

  committed
  1893881

  Copy full SHA for 1893881

• Merge pull request jax-ml#25831 from jax-ml:avoid-float0-tracers

  Show description for 4f106b8
  Google-ML-Automation

  committed
  4f106b8

  Copy full SHA for 4f106b8

• Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.

  Show description for c1de7c7
  

  dfm

  authored and

  Google-ML-Automation

  committed
  c1de7c7

  Copy full SHA for c1de7c7

• Activate FFI implementation of tridiagonal reduction on GPU.

  Show description for 39ce791
  

  dfm

  authored and

  Google-ML-Automation

  committed
  39ce791

  Copy full SHA for 39ce791

• Add a JVP rule for lax.linalg.tridiagonal_solve + various fixes.

  dfm

  committed
  167a48f

  Copy full SHA for 167a48f

• Simplify failure reporting for nightly CI job

  jakevdp

  committed
  7f55ccc

  Copy full SHA for 7f55ccc

• Update XLA dependency to use revision

  Show description for 1cc07dd
  Google-ML-Automation

  committed
  1cc07dd

  Copy full SHA for 1cc07dd

• jnp.linalg.solve: finalize deprecation of batched 1D solves

  jakevdp

  committed
  051abaf

  Copy full SHA for 051abaf

• Merge pull request jax-ml#25838 from jakevdp:simplify-nightly-reporting

  Show description for f7b6dee
  Google-ML-Automation

  committed
  f7b6dee

  Copy full SHA for f7b6dee

• Merge pull request jax-ml#25787 from dfm:tri-diag-jvp

  Show description for 016fca7
  Google-ML-Automation

  committed
  016fca7

  Copy full SHA for 016fca7

• Annotate several tests as thread-unsafe.

  Show description for 8f2f4b4
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  8f2f4b4

  Copy full SHA for 8f2f4b4

• Merge pull request jax-ml#25791 from mattjj:logsumexp-where-grad-nan

  Show description for aed7970
  Google-ML-Automation

  committed
  aed7970

  Copy full SHA for aed7970

• Merge pull request jax-ml#25741 from jakevdp:solve-dep

  Show description for 5d0ee43
  Google-ML-Automation

  committed
  5d0ee43

  Copy full SHA for 5d0ee43

• [internal] add deprecation test utilities

  jakevdp

  committed
  1ee0156

  Copy full SHA for 1ee0156

• Merge pull request jax-ml#25817 from jakevdp:deprecation-utils

  Show description for 743872d
  Google-ML-Automation

  committed
  743872d

  Copy full SHA for 743872d

• [Mosaic][TPU] Add a compatibility mode to Mosaic's canonicalization pass, skipping over elementwise and matmul op insertions and/or type compat casts.

  Show description for a16fbff
  Google-ML-Automation

  committed
  a16fbff

  Copy full SHA for a16fbff

• [Mosaic GPU] Enable loop carries in the pipeline emitter.

  Show description for 73b64b8
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  73b64b8

  Copy full SHA for 73b64b8

• [Mosaic GPU] Allow multiple gmem indexers on copies.

  Show description for 8e86bed
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  8e86bed

  Copy full SHA for 8e86bed

• [sharding_in_types] Introduce auto_mode, user_mode, auto_mode_ctx and user_mode_ctx as **private** APIs to make writing auto/user sharding in types code way easier and noise-free.

  Show description for a817f53
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  a817f53

  Copy full SHA for a817f53

Commits on Jan 11, 2025

• Fix pmap with sharded typed prng key

  Show description for aeac6b0
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  aeac6b0

  Copy full SHA for aeac6b0

• Update XLA dependency to use revision

  Show description for 4e72931
  Google-ML-Automation

  committed
  4e72931

  Copy full SHA for 4e72931

Commits on Jan 12, 2025

• [shardy] Fix cases in shardy where you have a nullary function with partially specified out_shardings (i.e. some out_sharding's are None and others are NamedShardings).

  Show description for 6b253b2
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6b253b2

  Copy full SHA for 6b253b2

• Update XLA dependency to use revision

  Show description for ed349e4
  Google-ML-Automation

  committed
  ed349e4

  Copy full SHA for ed349e4

Commits on Jan 13, 2025

• [Mosaic TPU][NFC] Remove redundant num_subelems attribute from CreateSubelementMaskOp

  Show description for 0930289
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  0930289

  Copy full SHA for 0930289

• [Mosaic TPU] Enable non-sublane-aligned bf16 2D load/stores for earlier TPU gens

  Show description for 7852045
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  7852045

  Copy full SHA for 7852045

• Use thread-safe initialization of LAPACK kernels.

  Show description for 91ffb64
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  91ffb64

  Copy full SHA for 91ffb64

• Add JAX unit test for Shardy which causes the compiler to introduce the mlir::tensor::TensorDialect. This was causing the compiler to crash.

  Show description for c14e5b4
  

  bartchr808

  authored and

  Google-ML-Automation

  committed
  c14e5b4

  Copy full SHA for c14e5b4

• [better_errors] Ensure that tracer errors in for_loop points to use code

  Show description for 3faff78
  gnecula

  committed
  3faff78

  Copy full SHA for 3faff78

• [Mosaic TPU] Add support for arith.fptosi with non-32bit source and target types

  Show description for 391bad8
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  391bad8

  Copy full SHA for 391bad8

• [Pallas TPU] Skip cast test incompatible with older libtpu builds

  Show description for aa51f2a
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  aa51f2a

  Copy full SHA for aa51f2a

• Update XLA dependency to use revision

  Show description for 9ee25e7
  Google-ML-Automation

  committed
  9ee25e7

  Copy full SHA for 9ee25e7

• Internal change.

  Show description for e72c148
  

  dfm

  authored and

  Google-ML-Automation

  committed
  e72c148

  Copy full SHA for e72c148

• Even more linearize fixes

  dougalm

  committed
  96769f9

  Copy full SHA for 96769f9

• Merge pull request jax-ml#25833 from jax-ml:more-linearize-fixes

  Show description for dabe27b
  Google-ML-Automation

  committed
  dabe27b

  Copy full SHA for dabe27b

• [Mosaic GPU] Fix layout API bugs.

  Show description for f69592a
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  f69592a

  Copy full SHA for f69592a

• [Pallas TPU] Add vector support to pl.debug_print

  Show description for 9ba1fd2
  

  ayaka14732

  authored and

  Google-ML-Automation

  committed
  9ba1fd2

  Copy full SHA for 9ba1fd2

Commits on Jan 14, 2025

• enable partitionable threefry by default

  Show description for a60ead6
  

  froystig

  authored and

  Google-ML-Automation

  committed
  a60ead6

  Copy full SHA for a60ead6

• Merge pull request jax-ml#25798 from gnecula:fix_fori_error

  Show description for 4f2f5fa
  Google-ML-Automation

  committed
  4f2f5fa

  Copy full SHA for 4f2f5fa

• #sdy dynamically choose which custom_partitioning API to use based on the current

  Show description for 74e912c
  

  bartchr808

  authored and

  Google-ML-Automation

  committed
  74e912c

  Copy full SHA for 74e912c

• Mention expected tangent aval in error message, see jax-ml#25517.

  dougalm

  committed
  7d11d12

  Copy full SHA for 7d11d12

• Merge pull request jax-ml#25875 from jax-ml:issue-25517

  Show description for 29a3dde
  Google-ML-Automation

  committed
  29a3dde

  Copy full SHA for 29a3dde

• Pin actions/checkout to a commit

  Show description for a1bbad6
  

  nitins17

  authored and

  Google-ML-Automation

  committed
  a1bbad6

  Copy full SHA for a1bbad6

• [sharding_in_types] Handle ShapeDtypeStruct inputs with sharding_in_types by registering the sharding on the aval properly created by SDS in it's pytype_aval_mapping.

  Show description for c72ed26
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  c72ed26

  Copy full SHA for c72ed26

• Merge branch 'rocm-main' into ci-upstream-sync-87_1

  charleshofer

  committed
  23fcd11

  Copy full SHA for 23fcd11