Code

CI: 11/22/24 upstream sync#148

Code

Merged

charleshofer merged 155 commits intorocm-mainROCm/jax:rocm-mainfrom ci-upstream-sync-34_1ROCm/jax:ci-upstream-sync-34_1Copy head branch name to clipboardNov 22, 2024

+6,010-2,401

Commits

Commits on Oct 28, 2024

• Add a link to Intel plugin for JAX

  mini-goel

  committed
  6842848

  Copy full SHA for 6842848

Commits on Nov 8, 2024

• Add float8_e4m3 and float8_e3m4 types support

  

  superbobry

  authored and

  apivovarov

  committed
  78da9fa

  Copy full SHA for 78da9fa

Commits on Nov 14, 2024

• Move Control Flow text from Sharp Bits into its own tutorial.

  emilyfertig

  committed
  e6f6a8a

  Copy full SHA for e6f6a8a

• jnp.logaddexp2: simplify implementation

  jakevdp

  committed
  d823f17

  Copy full SHA for d823f17

• Update array-api-tests commit

  jakevdp

  committed
  d0f3666

  Copy full SHA for d0f3666

• Merge pull request jax-ml#24864 from jakevdp:logaddexp2

  Show description for 303b792
  Google-ML-Automation

  committed
  303b792

  Copy full SHA for 303b792

• Add numpy.put_along_axis.

  carlosgmartin

  committed
  1f114b1

  Copy full SHA for 1f114b1

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

  Show description for 04d339d
  Google-ML-Automation

  committed
  04d339d

  Copy full SHA for 04d339d

• cleanup: delete unused argument from internal reduction helper

  jakevdp

  committed
  4a3e115

  Copy full SHA for 4a3e115

• Update XLA dependency to use revision

  Show description for c40d405
  Google-ML-Automation

  committed
  c40d405

  Copy full SHA for c40d405

• Add shard map replication rule for ffi_call.

  dfm

  committed
  41a0493

  Copy full SHA for 41a0493

• Merge pull request jax-ml#24567 from Intel-tensorflow:minigoel/intel-plugin

  Show description for 8e29212
  Google-ML-Automation

  committed
  8e29212

  Copy full SHA for 8e29212

• Merge pull request jax-ml#24905 from jakevdp:old-arg

  Show description for 5764afb
  Google-ML-Automation

  committed
  5764afb

  Copy full SHA for 5764afb

Commits on Nov 15, 2024

• Merge pull request jax-ml#24871 from carlosgmartin:numpy_put_along_axis

  Show description for 4fe9164
  Google-ML-Automation

  committed
  4fe9164

  Copy full SHA for 4fe9164

• Update array-api-tests commit

  jakevdp

  committed
  a115b2c

  Copy full SHA for a115b2c

• Merge pull request jax-ml#24881 from dfm:ffi-call-rep-rule

  Show description for c6051b3
  Google-ML-Automation

  committed
  c6051b3

  Copy full SHA for c6051b3

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

  Show description for 1c31860
  Google-ML-Automation

  committed
  1c31860

  Copy full SHA for 1c31860

• Merge pull request jax-ml#24862 from emilyfertig:emilyaf-control-flow-tutorial

  Show description for 4511f0c
  Google-ML-Automation

  committed
  4511f0c

  Copy full SHA for 4511f0c

• [sharding_in_types] Handle collective axes in lowering rules more generally. If any axis is collective, set all dims of aval to unspecified dims in wrap_with_sharding_op.

  Show description for 9a0e9e5
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9a0e9e5

  Copy full SHA for 9a0e9e5

• Set __module__ attribute for objects in jax.numpy

  jakevdp

  committed
  f652b6a

  Copy full SHA for f652b6a

• [Mosaic TPU] Support 1D concat: set implicit_dim to kSecondMinor to treat 1D (N,) as (1, N) and then tile it as (1, 128)

  Show description for 1471702
  Google-ML-Automation

  committed
  1471702

  Copy full SHA for 1471702

• Lower threefry as an out-of-line MLIR function on TPU.

  Show description for 23e9142
  hawkinsp

  committed
  23e9142

  Copy full SHA for 23e9142

• Merge pull request jax-ml#24913 from hawkinsp:threefry

  Show description for d808500
  Google-ML-Automation

  committed
  d808500

  Copy full SHA for d808500

• Add missing functions to jax.numpy type interface

  jakevdp

  committed
  5f94284

  Copy full SHA for 5f94284

• Add an example on logical operators to the tutorial.

  emilyfertig

  committed
  5f1e3f5

  Copy full SHA for 5f1e3f5

• Update XLA dependency to use revision

  Show description for 1780ff2
  Google-ML-Automation

  committed
  1780ff2

  Copy full SHA for 1780ff2

• DOC: update main landing page style

  Show description for 81cdc88
  

  barnesjoseph

  and

  jakevdp

  committed
  81cdc88

  Copy full SHA for 81cdc88

• Merge pull request jax-ml#24918 from emilyfertig:emilyaf-logical-op-example

  Show description for 605c605
  Google-ML-Automation

  committed
  605c605

  Copy full SHA for 605c605

• Merge pull request jax-ml#24914 from jakevdp:fix-pyi

  Show description for 1aa5de6
  Google-ML-Automation

  committed
  1aa5de6

  Copy full SHA for 1aa5de6

• Consolidate material on PRNGs and add a short summary to Key Concepts.

  emilyfertig

  committed
  225a2a5

  Copy full SHA for 225a2a5

• Merge pull request jax-ml#24917 from emilyfertig:emilyaf-sharp-bits

  Show description for efd2327
  Google-ML-Automation

  committed
  efd2327

  Copy full SHA for efd2327

Commits on Nov 16, 2024

• [sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot

  Show description for 8525ef2
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  8525ef2

  Copy full SHA for 8525ef2

• Adds a flag to control proxy env checking.

  Show description for 609dfac
  yliu120

  committed
  609dfac

  Copy full SHA for 609dfac

• Deduplicate constants in StableHLO lowering.

  Show description for 626aea0
  hawkinsp

  committed
  626aea0

  Copy full SHA for 626aea0

• Return a ndarray in shape_as_value if the shape is known to be constant.

  hawkinsp

  committed
  1d519f4

  Copy full SHA for 1d519f4

• Update XLA dependency to use revision

  Show description for 7b9914d
  Google-ML-Automation

  committed
  7b9914d

  Copy full SHA for 7b9914d

• Use a direct StableHLO lowering for pow.

  Show description for 8a6c560
  hawkinsp

  committed
  8a6c560

  Copy full SHA for 8a6c560

• Adds an env that can let users provide a custom version suffix for jax dev build.

  Show description for 27bf80a
  yliu120

  committed
  27bf80a

  Copy full SHA for 27bf80a

Commits on Nov 17, 2024

• Update XLA dependency to use revision

  Show description for 742cabc
  Google-ML-Automation

  committed
  742cabc

  Copy full SHA for 742cabc

Commits on Nov 18, 2024

• [AutoPGLE] Temporary disable pgle_test in the OSS.

  Show description for ed250b8
  Google-ML-Automation

  committed
  ed250b8

  Copy full SHA for ed250b8

• Merge pull request jax-ml#24930 from hawkinsp:dedup

  Show description for f7ae0f9
  Google-ML-Automation

  committed
  f7ae0f9

  Copy full SHA for f7ae0f9

• Merge pull request jax-ml#24933 from hawkinsp:pow

  Show description for afdc792
  Google-ML-Automation

  committed
  afdc792

  Copy full SHA for afdc792

• Add a GPU implementation of lax.linalg.eig.

  Show description for ccb3317
  

  dfm

  authored and

  Google-ML-Automation

  committed
  ccb3317

  Copy full SHA for ccb3317

• Merge pull request jax-ml#24932 from hawkinsp:gather

  Show description for 65f9c78
  Google-ML-Automation

  committed
  65f9c78

  Copy full SHA for 65f9c78

• Add new CI script for running Bazel GPU presubmits

  Show description for 1418739
  

  nitins17

  authored and

  Google-ML-Automation

  committed
  1418739

  Copy full SHA for 1418739

• Merge pull request jax-ml#24912 from jakevdp:jnp-module

  Show description for 05d66d7
  Google-ML-Automation

  committed
  05d66d7

  Copy full SHA for 05d66d7

• Merge pull request jax-ml#24916 from jakevdp:update-lp

  Show description for 2de40e7
  Google-ML-Automation

  committed
  2de40e7

  Copy full SHA for 2de40e7

• Make logaddexp and logaddexp2 into ufuncs

  jakevdp

  committed
  e9864c6

  Copy full SHA for e9864c6

• Merge pull request jax-ml#24903 from jakevdp:logsumexp

  Show description for 297a4e5
  Google-ML-Automation

  committed
  297a4e5

  Copy full SHA for 297a4e5

• Return SingleDeviceSharding instead of GSPMDShardings when there is only 1 device during compiled.input_shardings call.

  Show description for 6fe7b17
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6fe7b17

  Copy full SHA for 6fe7b17

• fix typo in numpy/__init__.pyi

  jakevdp

  committed
  5bebd0f

  Copy full SHA for 5bebd0f

• Merge pull request jax-ml#24952 from jakevdp:fix-pyi

  Show description for 70b05f6
  Google-ML-Automation

  committed
  70b05f6

  Copy full SHA for 70b05f6

• [SDY] fix JAX layouts tests for Shardy.

  Show description for 0ed6eae
  

  Varcho

  authored and

  Google-ML-Automation

  committed
  0ed6eae

  Copy full SHA for 0ed6eae

• Disable some complex function accuracy tests that fail on Mac ARM.

  hawkinsp

  committed
  461a250

  Copy full SHA for 461a250

• Filter custom dtypes by supported_dtypes in _LazyDtypes.

  Show description for f325051
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  f325051

  Copy full SHA for f325051

• [Pallas] Increase test coverage of pl.dot.

  Show description for a60ef6e
  

  WindQAQ

  authored and

  Google-ML-Automation

  committed
  a60ef6e

  Copy full SHA for a60ef6e

• Merge pull request jax-ml#24957 from hawkinsp:arm

  Show description for 16ed283
  Google-ML-Automation

  committed
  16ed283

  Copy full SHA for 16ed283

• Update XLA dependency to use revision

  Show description for b3ca6c4
  Google-ML-Automation

  committed
  b3ca6c4

  Copy full SHA for b3ca6c4

• Merge pull request jax-ml#23585 from apivovarov:float8_e4m3

  Show description for 91891cb
  Google-ML-Automation

  committed
  91891cb

  Copy full SHA for 91891cb

• Adds font fallbacks

  barnesjoseph

  committed
  d4316b5

  Copy full SHA for d4316b5

• Merge pull request jax-ml#24958 from barnesjoseph:add-font-fallback

  Show description for 6952ddf
  Google-ML-Automation

  committed
  6952ddf

  Copy full SHA for 6952ddf

• Delete _normalized_spec from NamedSharding

  Show description for e904c17
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  e904c17

  Copy full SHA for e904c17

Commits on Nov 19, 2024

• Fix a bug where mesh checking was not correct

  Show description for 2c68569
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  2c68569

  Copy full SHA for 2c68569

• [pallas] Minor simplifications to Pallas interpreter.

  Show description for 45c9c0a
  

  chr1sj0nes

  authored and

  Google-ML-Automation

  committed
  45c9c0a

  Copy full SHA for 45c9c0a

• Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.

  Show description for c5e8ae8
  hawkinsp

  committed
  c5e8ae8

  Copy full SHA for c5e8ae8

• Merge pull request jax-ml#24945 from hawkinsp:gamma

  Show description for 4a9346e
  Google-ML-Automation

  committed
  4a9346e

  Copy full SHA for 4a9346e

• Merge pull request jax-ml#24861 from yliu120:add_versions

  Show description for 58103e5
  Google-ML-Automation

  committed
  58103e5

  Copy full SHA for 58103e5

• [Mosaic TPU] Support relayout for mask vector

  Show description for 0fe77bc
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  0fe77bc

  Copy full SHA for 0fe77bc

• Merge pull request jax-ml#24853 from yliu120:check_proxy_envs

  Show description for 12a43f1
  Google-ML-Automation

  committed
  12a43f1

  Copy full SHA for 12a43f1

• Implement lax.pad in Pallas.

  Show description for d397dd9
  Google-ML-Automation

  committed
  d397dd9

  Copy full SHA for d397dd9

• [AutoPGLE] Use compile options to override debug options instead of XLA_FLAGS.

  Show description for da50ad7
  Google-ML-Automation

  committed
  da50ad7

  Copy full SHA for da50ad7

• Fix some typos

  nireekshak

  committed
  1458d3d

  Copy full SHA for 1458d3d

• fix(docs): typos in macro name

  Show description for d912034
  jeertmans

  committed
  d912034

  Copy full SHA for d912034

• Merge pull request jax-ml#24942 from jeertmans:patch-1

  Show description for 9d3eda1
  Google-ML-Automation

  committed
  9d3eda1

  Copy full SHA for 9d3eda1

• Merge pull request jax-ml#24968 from nireekshak:testingbranch

  Show description for 6929a97
  Google-ML-Automation

  committed
  6929a97

  Copy full SHA for 6929a97

• Add missing version guard in GPU tests for jnp.poly.

  Show description for 3556a83
  

  dfm

  authored and

  Google-ML-Automation

  committed
  3556a83

  Copy full SHA for 3556a83

• [Mosaic TPU] Add general tpu.vector_store and support masked store.

  Show description for 6c31efa
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  6c31efa

  Copy full SHA for 6c31efa

• Add alternate implementation of threefry as a pallas kernel.

  Show description for c44f11d
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  c44f11d

  Copy full SHA for c44f11d

• Add test utility for accessing jaxlib version tuple.

  Show description for a59bbb7
  

  dfm

  authored and

  Google-ML-Automation

  committed
  a59bbb7

  Copy full SHA for a59bbb7

• Add a new API jax.lax.split.

  Show description for 2c80d1a
  hawkinsp

  committed
  2c80d1a

  Copy full SHA for 2c80d1a

• Merge pull request jax-ml#24970 from hawkinsp:split

  Show description for 2075b09
  Google-ML-Automation

  committed
  2075b09

  Copy full SHA for 2075b09

• [pallas:mosaic_gpu] copy_gmem_to_smem no longer requires barrier to be a keyword argument

  Show description for 1bf70fb
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  1bf70fb

  Copy full SHA for 1bf70fb

• [Mosaic] Add target core type parameter to tpu.sem_signal

  Show description for 0d36b0b
  

  naummo

  authored and

  Google-ML-Automation

  committed
  0d36b0b

  Copy full SHA for 0d36b0b

• Update XLA dependency to use revision

  Show description for 3161a28
  Google-ML-Automation

  committed
  3161a28

  Copy full SHA for 3161a28

• Move JAX example to public XLA:CPU API

  Show description for 42fbd30
  

  changm

  authored and

  Google-ML-Automation

  committed
  42fbd30

  Copy full SHA for 42fbd30

• Reverts 2075b09

  Show description for 525b646
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  525b646

  Copy full SHA for 525b646

• [Mosaic] Extend tpu.sem_signal with subcore_id

  Show description for c04aec9
  

  naummo

  authored and

  Google-ML-Automation

  committed
  c04aec9

  Copy full SHA for c04aec9

Commits on Nov 20, 2024

• Make deprecated jax.experimental.array_api module visibility internal-only

  Show description for 8c71d1a
  

  Jake VanderPlas

  authored and

  Google-ML-Automation

  committed
  8c71d1a

  Copy full SHA for 8c71d1a

• Fix a bug where constant deduplication used an inappropriate inequality.

  Show description for 867a361
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  867a361

  Copy full SHA for 867a361

• [Mosaic] Add tpu.log verification on SC

  Show description for 6c291d6
  

  naummo

  authored and

  Google-ML-Automation

  committed
  6c291d6

  Copy full SHA for 6c291d6

• represent random.key_impl of builtin RNGs by canonical string name

  Show description for 4bb8107
  froystig

  committed
  4bb8107

  Copy full SHA for 4bb8107

• Add test_compute_on_host_shared_sharding in memories_test

  Show description for 4d60db1
  Google-ML-Automation

  committed
  4d60db1

  Copy full SHA for 4d60db1

• Merge pull request jax-ml#24593 from froystig:random-dtypes

  Show description for ae46b75
  Google-ML-Automation

  committed
  ae46b75

  Copy full SHA for ae46b75

• [mosaic_gpu] Fix signedness handling in FragmentedArray._pointwise.

  Show description for 1afb05e
  

  petebu

  authored and

  Google-ML-Automation

  committed
  1afb05e

  Copy full SHA for 1afb05e

• [pallas:mosaic_gpu] Add Pallas Mosaic GPU lowering for jax.lax.bitcast_convert_type.

  Show description for 14da7eb
  

  petebu

  authored and

  Google-ML-Automation

  committed
  14da7eb

  Copy full SHA for 14da7eb

• [pallas:mosaic_gpu] copy_smem_to_gmem now supports wait_read_only

  Show description for c76e5fe
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  c76e5fe

  Copy full SHA for c76e5fe

• [mosaic_gpu] Fixed FragmentedArray comparisons with literals

  Show description for f442d40
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  f442d40

  Copy full SHA for f442d40

• [mosaic_gpu] Handle older jaxlibs in the profiler module

  Show description for 04e4c69
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  04e4c69

  Copy full SHA for 04e4c69

• [pallas] Do not skip vmap tests on GPU when x64 is enabled

  Show description for 1df4b5f
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  1df4b5f

  Copy full SHA for 1df4b5f

• Update XLA dependency to use revision

  Show description for a582df0
  Google-ML-Automation

  committed
  a582df0

  Copy full SHA for a582df0

• Mention python 3.13 in docs & package metadata

  jakevdp

  committed
  a4266b5

  Copy full SHA for a4266b5

• Simplify handling of DotAlgorithmPreset output types.

  Show description for 1e9e85a
  

  chr1sj0nes

  authored and

  Google-ML-Automation

  committed
  1e9e85a

  Copy full SHA for 1e9e85a

• Deprecate several private APIs in jax.lib

  jakevdp

  committed
  85e2969

  Copy full SHA for 85e2969

• Merge pull request jax-ml#25007 from jakevdp:deps

  Show description for 800add2
  Google-ML-Automation

  committed
  800add2

  Copy full SHA for 800add2

• Merge pull request jax-ml#25005 from jakevdp:py313

  Show description for 439d34d
  Google-ML-Automation

  committed
  439d34d

  Copy full SHA for 439d34d

• Fix KeyError recently introduced in cloud_tpu_init.py

  Show description for 6222592
  skye

  committed
  6222592

  Copy full SHA for 6222592

• [pallas mgpu] Lowering for while loops as long as they are secretly for loops.

  Show description for 8d84f28
  

  cperivol

  authored and

  Google-ML-Automation

  committed
  8d84f28

  Copy full SHA for 8d84f28

• Make a direct linearize trace.

  Show description for d0f17c0
  dougalm

  committed
  d0f17c0

  Copy full SHA for d0f17c0

• Merge pull request jax-ml#25004 from jax-ml:linearize-trace

  Show description for eab9026
  Google-ML-Automation

  committed
  eab9026

  Copy full SHA for eab9026

• Remove internal KeyArray alias

  Show description for fee272e
  jakevdp

  committed
  fee272e

  Copy full SHA for fee272e

• Don't psum over auto mesh dims in _unmentioned2.

  Show description for 2c9b917
  

  pschuh

  authored and

  Google-ML-Automation

  committed
  2c9b917

  Copy full SHA for 2c9b917

• [pallas:mosaic_gpu] Avoid using multiple indexers in the parallel grid test

  Show description for 9584ee3
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  9584ee3

  Copy full SHA for 9584ee3

• Set __module__ attribute of jax.numpy.linalg APIs

  jakevdp

  committed
  621e39d

  Copy full SHA for 621e39d

• Merge pull request jax-ml#25008 from skye:barrier

  Show description for 1a3e693
  Google-ML-Automation

  committed
  1a3e693

  Copy full SHA for 1a3e693

• Mention stackless in the release notes.

  hawkinsp

  committed
  dfe27a1

  Copy full SHA for dfe27a1

• Merge pull request jax-ml#25013 from hawkinsp:relnotes

  Show description for 19b4996
  Google-ML-Automation

  committed
  19b4996

  Copy full SHA for 19b4996

• [sharding_in_types] Make flash_attention forward pass in TPU pallas work nicely with sharding in types. Backward pass is still busted which I will fix in follow up CLs.

  Show description for 40fc659
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  40fc659

  Copy full SHA for 40fc659

• [Pallas TPU] Support masked store

  Show description for 9d2f62f
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  9d2f62f

  Copy full SHA for 9d2f62f

• Merge pull request jax-ml#25011 from jakevdp:jnp-linalg-module

  Show description for d219439
  Google-ML-Automation

  committed
  d219439

  Copy full SHA for d219439

• [sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

  Show description for 9b94180
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9b94180

  Copy full SHA for 9b94180

• Update XLA dependency to use revision

  Show description for 6fe7804
  Google-ML-Automation

  committed
  6fe7804

  Copy full SHA for 6fe7804

• [array api] use most recent version of array_api_tests

  jakevdp

  committed
  f749fca

  Copy full SHA for f749fca

• DOC: add examples for jax.lax.pad

  jakevdp

  committed
  2699e95

  Copy full SHA for 2699e95

Commits on Nov 21, 2024

• Merge pull request jax-ml#25019 from jakevdp:lax-pad-doc

  Show description for 334bd4d
  Google-ML-Automation

  committed
  334bd4d

  Copy full SHA for 334bd4d

• jax.lax.pad: improve input validation

  jakevdp

  committed
  1782588

  Copy full SHA for 1782588

• Adds Google Sans font

  barnesjoseph

  committed
  bf7f9aa

  Copy full SHA for bf7f9aa

• Merge pull request jax-ml#25020 from jakevdp:lax-pad-validation

  Show description for f39392e
  Google-ML-Automation

  committed
  f39392e

  Copy full SHA for f39392e

• [Pallas] Use Pallas cost estimator for flash attention.

  Show description for 1f6152d
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  1f6152d

  Copy full SHA for 1f6152d

• [sharding_in_types] Add pad_p support to sharding_in_types to handle transpose to slice correctly.

  Show description for 840cf3f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  840cf3f

  Copy full SHA for 840cf3f

• [Mosaic TPU] Add bound check for general vector store op.

  Show description for 869a533
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  869a533

  Copy full SHA for 869a533

• [sharding_in_types] Add concatenate_p support

  Show description for 6568713
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6568713

  Copy full SHA for 6568713

• Reverts c04aec9

  Show description for e72b449
  

  naummo

  authored and

  Google-ML-Automation

  committed
  e72b449

  Copy full SHA for e72b449

• [pallas:mosaic_gpu] Pulled delay_release into emit_pipeline

  Show description for f18df8f
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  f18df8f

  Copy full SHA for f18df8f

• Integrate LLVM at llvm/llvm-project@33fcd6acc755

  Show description for 1bc9df4
  

  metaflow

  authored and

  Google-ML-Automation

  committed
  1bc9df4

  Copy full SHA for 1bc9df4

• [shape_poly] Adding shape polymorphism support for the state primitives.

  gnecula

  committed
  0831e2e

  Copy full SHA for 0831e2e

• Run the TPU workflow on new self-hosted runners

  Show description for 7d7a0fa
  

  nitins17

  authored and

  Google-ML-Automation

  committed
  7d7a0fa

  Copy full SHA for 7d7a0fa

• [JAX] Ignore xla_gpu_experimental_autotune_cache_mode when calculating module hash.

  Show description for bf0150b
  Google-ML-Automation

  committed
  bf0150b

  Copy full SHA for bf0150b

• Merge pull request jax-ml#25015 from barnesjoseph:add-google-sans

  Show description for 7335267
  Google-ML-Automation

  committed
  7335267

  Copy full SHA for 7335267

• Fix cron schedule to run past minute 0 every 2nd hour

  Show description for 1e6654a
  

  nitins17

  authored and

  Google-ML-Automation

  committed
  1e6654a

  Copy full SHA for 1e6654a

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

  Show description for b1b1ad6
  Google-ML-Automation

  committed
  b1b1ad6

  Copy full SHA for b1b1ad6

• [mgpu] Pointwise op can handle LHS splats.

  Show description for 1d2dc17
  

  cperivol

  authored and

  Google-ML-Automation

  committed
  1d2dc17

  Copy full SHA for 1d2dc17

• [pallas] Add more test cases for Triton bitcast_convert_type lowering rule.

  Show description for 2178ed2
  

  petebu

  authored and

  Google-ML-Automation

  committed
  2178ed2

  Copy full SHA for 2178ed2

• Merge pull request jax-ml#25034 from gnecula:poly_state

  Show description for e707ede
  Google-ML-Automation

  committed
  e707ede

  Copy full SHA for e707ede

• Fix false positive debug_nans error caused by NaNs that are properly handled in jax.scipy.stats.gamma

  Show description for 96c0129
  

  dfm

  authored and

  Google-ML-Automation

  committed
  96c0129

  Copy full SHA for 96c0129

• [pallas:mosaic_gpu] emit_pipeline now correctly supports BlockSpecs in GMEM

  Show description for 1efef6b
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  1efef6b

  Copy full SHA for 1efef6b

• Merge pull request jax-ml#25009 from jakevdp:keyarray

  Show description for c1ae13b
  Google-ML-Automation

  committed
  c1ae13b

  Copy full SHA for c1ae13b

• Remove unneeded dependency from rocm_plugin_extension.

  Show description for f3e7e68
  

  klucke

  authored and

  Google-ML-Automation

  committed
  f3e7e68

  Copy full SHA for f3e7e68

• [Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.

  Show description for f899d51
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  f899d51

  Copy full SHA for f899d51

• Update XLA dependency to use revision

  Show description for 26443bb
  Google-ML-Automation

  committed
  26443bb

  Copy full SHA for 26443bb

• [Pallas] Add readme page for debugging tips.

  Show description for 344d0d9
  

  justinjfu

  authored and

  Google-ML-Automation

  committed
  344d0d9

  Copy full SHA for 344d0d9

Commits on Nov 22, 2024

• Change signature of linearization rules.

  Show description for 170718c
  dougalm

  committed
  170718c

  Copy full SHA for 170718c

• Merge pull request jax-ml#25048 from jax-ml:linearization-rule-signature

  Show description for 3d79df2
  Google-ML-Automation

  committed
  3d79df2

  Copy full SHA for 3d79df2

• [sharding_in_types] Add scan support to sharding_in_types. There are a couple of changes here

  Show description for 355589f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  355589f

  Copy full SHA for 355589f

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

  charleshofer

  committed
  a450bb0

  Copy full SHA for a450bb0

• Longer timeout for doc render

  charleshofer

  committed
  846697f

  Copy full SHA for 846697f