Commits

Branch selector

 c78ca04

User selector

All users

All time

Commit History

Commits on Dec 19, 2024

• Add experimental support for building JAX CPU and GPU wheels with GCC.

  Show description for c78ca04
  Google-ML-Automation

  committed
  c78ca04

  Copy full SHA for c78ca04

• Merge pull request jax-ml#25516 from jakevdp:jep-effver

  Show description for 49c9246
  Google-ML-Automation

  committed
  49c9246

  Copy full SHA for 49c9246

• Merge pull request jax-ml#25590 from jakevdp:fix-one-hot-float

  Show description for 8402a98
  Google-ML-Automation

  committed
  8402a98

  Copy full SHA for 8402a98

• Always suppress the differing_executors Hypothesis health check

  Show description for 23000a3
  

  apaszke

  authored and

  Google-ML-Automation

  committed
  23000a3

  Copy full SHA for 23000a3

• [Mosaic GPU] Commit to using Vectors everywhere (and no Tensors).

  Show description for 3915f4a
  

  bchetioui

  authored and

  Google-ML-Automation

  committed
  3915f4a

  Copy full SHA for 3915f4a

• [jaxlib] Added a missing pytype_dep to :_triton_ext

  Show description for c4fae4a
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  c4fae4a

  Copy full SHA for c4fae4a

• [pallas:triton] Fixed a typo in a type annotation

  Show description for af7a31f
  

  superbobry

  authored and

  Google-ML-Automation

  committed
  af7a31f

  Copy full SHA for af7a31f

• jax.nn.one_hot: deprecate non-integer inputs

  jakevdp

  committed
  8c3c441

  Copy full SHA for 8c3c441

• 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

Pagination

Previous

Next