Commits

Branch selector

 main

User selector

All users

Datepicker

All time

Commit History

Commits on Mar 28, 2024

• [XLA:GPU] Deprecate Triton codegen before Ampere.

  Show description for 1e1906b
  

  bchetioui

  authored and

  jax authors

  committed
  1e1906b

  Copy full SHA for 1e1906b

• Update XLA dependency to use revision

  Show description for 3afbeed
  

  jax authors

  committed
  3afbeed

  Copy full SHA for 3afbeed

• [XLA:Mosaic] Remove duplicate headers in debug assert insertion.

  Show description for 67f4f60
  

  bythew3i

  authored and

  jax authors

  committed
  67f4f60

  Copy full SHA for 67f4f60

• Add custom call on output along with S(5) because XLA requires the custom call to show the transfer.

  Show description for 9e86aa5
  

  yashk2810

  authored and

  jax authors

  committed
  9e86aa5

  Copy full SHA for 9e86aa5

Commits on Mar 27, 2024

• Do not deadlock the GPU if a pure_callback dispatches a GPU kernel

  Show description for ec73c40
  

  superbobry

  authored and

  jax authors

  committed
  ec73c40

  Copy full SHA for ec73c40

• Supports PjRt CPU array converting py array when the CPU PjRt arrays have non-default layouts.

  Show description for 7f30ab5
  

  Yunlong Liu

  authored and

  jax authors

  committed
  7f30ab5

  Copy full SHA for 7f30ab5

• Allow allow_spmd_propagation_to_output to be generated for outputs annotated with pjit.AUTO

  Show description for 66877c9
  

  jax authors

  committed
  66877c9

  Copy full SHA for 66877c9

• [export] Increase minimum serialization version to 9.

  Show description for c0c918a
  

  gnecula

  authored and

  jax authors

  committed
  c0c918a

  Copy full SHA for c0c918a

• Fix some load orderings for buildifier

  Show description for 023930d
  

  MichaelHudgins

  authored and

  jax authors

  committed
  023930d

  Copy full SHA for 023930d

• Update XLA dependency to use revision

  Show description for b480fd5
  

  jax authors

  committed
  b480fd5

  Copy full SHA for b480fd5

• Reverts 0dde8f7

  Show description for fa9f02b
  

  mattjj

  authored and

  jax authors

  committed
  fa9f02b

  Copy full SHA for fa9f02b

• Merge pull request jax-ml#20445 from mattjj:scan-dont-traverse-body-jaxpr-in-lowering-2

  Show description for 0dde8f7
  

  jax authors

  committed
  0dde8f7

  Copy full SHA for 0dde8f7

• Guard host transfers inside pure_callbacks from deadlocking the TPU.

  Show description for 0b09762
  

  pschuh

  authored and

  jax authors

  committed
  0b09762

  Copy full SHA for 0b09762

• [scan] don't traverse body jaxpr in lowering

  Show description for 9474b46
  mattjj

  committed
  9474b46

  Copy full SHA for 9474b46

Commits on Mar 26, 2024

• Removed obsolete call to libtpu_module.configure_library_path().

  Show description for 4a9c8d1
  

  Jieying Luo

  authored and

  jax authors

  committed
  4a9c8d1

  Copy full SHA for 4a9c8d1

• Remove the canonicalization to GSPMDSharding internally in jit. This is not required anymore since the caches are split into tracing, lowering and compilation.

  Show description for 6e0c955
  

  yashk2810

  authored and

  jax authors

  committed
  6e0c955

  Copy full SHA for 6e0c955

• Add commentary and clean the pallas bidirectional collective allgather matmul test.

  Show description for f8d6669
  

  levskaya

  authored and

  jax authors

  committed
  f8d6669

  Copy full SHA for f8d6669

• Merge pull request jax-ml#20436 from pearu:pearu/expm1-2

  Show description for 76af8be
  

  jax authors

  committed
  76af8be

  Copy full SHA for 76af8be

• Update complex function accurancy tests for expm1

  pearu

  committed
  c82f061

  Copy full SHA for c82f061

• Removed double-printing of TTIR in Pallas GPU lowering

  Show description for 18c885d
  

  superbobry

  authored and

  jax authors

  committed
  18c885d

  Copy full SHA for 18c885d

• Fix the pjit test failing on v5e

  Show description for c78054d
  

  yashk2810

  authored and

  jax authors

  committed
  c78054d

  Copy full SHA for c78054d

• Skip fused_attention_stablehlo_test.py.

  Show description for 6afa523
  

  jax authors

  committed
  6afa523

  Copy full SHA for 6afa523

• [callback] Fix io_callback for callbacks that return Python literals.

  Show description for 75db481
  

  gnecula

  authored and

  jax authors

  committed
  75db481

  Copy full SHA for 75db481

• [XLA:GPU] Add option to return FDO profile as textproto.

  Show description for 33cf53c
  

  cheshire

  authored and

  jax authors

  committed
  33cf53c

  Copy full SHA for 33cf53c

• Update XLA dependency to use revision

  Show description for 86a7086
  

  jax authors

  committed
  86a7086

  Copy full SHA for 86a7086

• Enable extra args with input output aliasing

  Show description for f93c320
  

  sharadmv

  authored and

  jax authors

  committed
  f93c320

  Copy full SHA for f93c320

• Use the information in allow_spmd_sharding_propagation_to_output and allow_spmd_sharding_propagation_to_parameters to determine what input and output tuple elements we are allowed to modfy the shar…

  Show description for 69980a2
  

  jax authors

  committed
  69980a2

  Copy full SHA for 69980a2

Commits on Mar 25, 2024

• Merge pull request jax-ml#20257 from Cjkkkk:sdpa_training

  Show description for c724eab
  

  jax authors

  committed
  c724eab

  Copy full SHA for c724eab

• move checks to setup

  Cjkkkk

  committed
  f51d80e

  Copy full SHA for f51d80e

• Avoid jax_explain_cache_misses unpacking error.

  Show description for e3bbd67
  

  jax authors

  committed
  e3bbd67

  Copy full SHA for e3bbd67

• Remove support for CUDA 11.

  Show description for 0be07e6
  

  jax authors

  committed
  0be07e6

  Copy full SHA for 0be07e6

• Reverts 2a4e1ca

  Show description for 19e6156
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  19e6156

  Copy full SHA for 19e6156

• add 4 gpus guard for inference

  Cjkkkk

  committed
  dacf009

  Copy full SHA for dacf009

• [Take 2] Expose .layout on jax.Array. Also add checks in the AOT path to make sure that the input Array's layout matches the layout given to jax.jit.

  Show description for 25d01e9
  

  yashk2810

  authored and

  jax authors

  committed
  25d01e9

  Copy full SHA for 25d01e9

• Pallas TPU now broadcasts operands when lowering bitwise ops

  Show description for b9e699f
  

  superbobry

  authored and

  jax authors

  committed
  b9e699f

  Copy full SHA for b9e699f

Pagination

Previous

Next