Commits

Branch selector

 5dc37d3

User selector

All users

Datepicker

All time

Commit History

Commits on Dec 19, 2024

• Remove internal uses of api_util.shaped_abstractify

  jakevdp

  committed
  5dc37d3

  Copy full SHA for 5dc37d3

• Merge pull request jax-ml#25595 from jakevdp:mv-shaped-abstractify

  Show description for 7680532
  Google-ML-Automation

  committed
  7680532

  Copy full SHA for 7680532

• [Mosaic TPU] Guard tests for new features by the libtpu version

  Show description for ad00ec1
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  ad00ec1

  Copy full SHA for ad00ec1

• [Pallas:TPU] Use self.pallas_call to properly handle interpret mode

  Show description for 4515949
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  4515949

  Copy full SHA for 4515949

• [Mosaic GPU] Add a new tiled layout, optimized for upcasting before WGMMA

  Show description for 006c65d
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  006c65d

  Copy full SHA for 006c65d

• [Mosaic GPU] Replace the dialect's layout enum with layouts holding the proper

  Show description for 66ad208
  

  bchetioui

  authored and

  Google-ML-Automation

  committed
  66ad208

  Copy full SHA for 66ad208

• [Mosaic:TPU][NFC] Small cleanup of extui rule in apply-vector-layout

  Show description for 8188c57
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  8188c57

  Copy full SHA for 8188c57

• [Mosaic GPU][NFC] Move up the definition of ThreadSemantics.

  Show description for 16f0f97
  

  dimitar-asenov

  authored and

  Google-ML-Automation

  committed
  16f0f97

  Copy full SHA for 16f0f97

• Add lax.composite primitive

  Show description for d206cc3
  

  ghpvnist

  authored and

  Google-ML-Automation

  committed
  d206cc3

  Copy full SHA for d206cc3

• Refactor: move shaped_abstractify to core

  jakevdp

  committed
  676070f

  Copy full SHA for 676070f

• [Pallas] Fix lowering tests for reduction ops

  Show description for 613a0cd
  

  ayaka14732

  authored and

  Google-ML-Automation

  committed
  613a0cd

  Copy full SHA for 613a0cd

• Merge pull request jax-ml#25593 from mattjj:ref-errors-4

  Show description for f65eced
  Google-ML-Automation

  committed
  f65eced

  Copy full SHA for f65eced

• Account for tokens in allow_spmd_sharding_propagation_to_parameters and allow_spmd_sharding_propagation_to_output compile options

  Show description for 9041b02
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9041b02

  Copy full SHA for 9041b02

• Merge pull request jax-ml#25596 from mattjj:whitespace

  Show description for 46b18d2
  Google-ML-Automation

  committed
  46b18d2

  Copy full SHA for 46b18d2

• add mutable array ref error checks to scan

  mattjj

  committed
  e528562

  Copy full SHA for e528562

• [Mosaic TPU] Support i32 vector multi reduction except cross lane.

  Show description for 3a5c4da
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  3a5c4da

  Copy full SHA for 3a5c4da

• [Mosaic] Improve error verbosity of tpu.memref_slice verification

  Show description for 6bcec91
  

  naummo

  authored and

  Google-ML-Automation

  committed
  6bcec91

  Copy full SHA for 6bcec91

Commits on Dec 18, 2024

• remove whitespace from distributed arrays doc

  mattjj

  committed
  c657682

  Copy full SHA for c657682

• Re-land changes from jax-ml#25555

  Show description for 89a54a9
  

  Jake VanderPlas

  authored and

  Google-ML-Automation

  committed
  89a54a9

  Copy full SHA for 89a54a9

• [Mosaic] Verify that the target IDs are provided in remote DMAs

  Show description for de359f5
  

  naummo

  authored and

  Google-ML-Automation

  committed
  de359f5

  Copy full SHA for de359f5

• [sharding_in_types] Check out_avals with mesh context too. This is because users can pass their own shardings to functions like einsum, reshape, broadcast`, etc

  Show description for af63e44
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  af63e44

  Copy full SHA for af63e44

• [Mosaic:TPU][NFC] Delete unused functions

  Show description for 13e721a
  

  tlongeri

  authored and

  Google-ML-Automation

  committed
  13e721a

  Copy full SHA for 13e721a

• [Pallas] Add version guard for non-32-bit selection in test and fix github build failure.

  Show description for 74eca13
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  74eca13

  Copy full SHA for 74eca13

• [Mosaic GPU][NFC] Split LaunchContext into a separate file

  Show description for 1ed5fa6
  

  dimitar-asenov

  authored and

  Google-ML-Automation

  committed
  1ed5fa6

  Copy full SHA for 1ed5fa6

• [Mosaic TPU] Support direct cast i8 vector to mask

  Show description for bf692ef
  

  bythew3i

  authored and

  Google-ML-Automation

  committed
  bf692ef

  Copy full SHA for bf692ef

• fix bug with jax.remat static_argnums not supporting int

  Show description for bb16d5a
  

  mattjj

  authored and

  Google-ML-Automation

  committed
  bb16d5a

  Copy full SHA for bb16d5a

• Update XLA dependency to use revision

  Show description for 53bff86
  Google-ML-Automation

  committed
  53bff86

  Copy full SHA for 53bff86

• Remove jax_enable_memories config flag. It defaulted to True for a very long time and it's time to remove the flag.

  Show description for 8b73480
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  8b73480

  Copy full SHA for 8b73480

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

  Show description for 8abb1a7
  Google-ML-Automation

  committed
  8abb1a7

  Copy full SHA for 8abb1a7

• [JAX] Fix a small bug if shardings is tuple.

  Show description for 72e5ca9
  

  maxwillzq

  authored and

  Google-ML-Automation

  committed
  72e5ca9

  Copy full SHA for 72e5ca9

• Merge pull request jax-ml#25569 from hawkinsp:numpyver

  Show description for 464e5a2
  Google-ML-Automation

  committed
  464e5a2

  Copy full SHA for 464e5a2

• More linearize fixes

  dougalm

  committed
  2be9a69

  Copy full SHA for 2be9a69

• Migrate mhlo dialect extension to nanobind.

  Show description for 3f24dfd
  

  hawkinsp

  authored and

  Google-ML-Automation

  committed
  3f24dfd

  Copy full SHA for 3f24dfd

• Partial discharge for scan_p ops.

  Show description for aaabb97
  

  cperivol

  authored and

  Google-ML-Automation

  committed
  aaabb97

  Copy full SHA for aaabb97

• Use capture_stderr instead of packing sys.stderr

  Show description for 2259a13
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  2259a13

  Copy full SHA for 2259a13

Pagination

Previous

Next