Commits

Branch selector

 main

User selector

All users

Datepicker

All time

Commit History

Commits on May 28, 2024

• [Mosaic] Add relayout for (1, 128 * packing) -> (packing, 128).

  Show description for 8b95853
  

  tlongeri

  authored and

  jax authors

  committed
  8b95853

  Copy full SHA for 8b95853

• Merge pull request jax-ml#21452 from jakevdp:take-fill-value

  Show description for 8fec5d6
  

  jax authors

  committed
  8fec5d6

  Copy full SHA for 8fec5d6

• Merge pull request jax-ml#21193 from sh0416:main

  Show description for 82ad1da
  

  jax authors

  committed
  82ad1da

  Copy full SHA for 82ad1da

• Merge pull request jax-ml#21453 from jakevdp:index-update-annotation

  Show description for 83871d3
  

  jax authors

  committed
  83871d3

  Copy full SHA for 83871d3

• Update the deprecation message of backend and device argument of jit to be more actionable.

  Show description for ff3db9b
  

  yashk2810

  authored and

  jax authors

  committed
  ff3db9b

  Copy full SHA for ff3db9b

• Merge pull request jax-ml#21451 from gnecula:poly_pmap

  Show description for de28ee6
  

  jax authors

  committed
  de28ee6

  Copy full SHA for de28ee6

• Fix type annotation for index-update

  jakevdp

  committed
  3ade023

  Copy full SHA for 3ade023

• jnp.take_along_axis: support fill_value

  jakevdp

  committed
  2d23a66

  Copy full SHA for 2d23a66

• Merge pull request jax-ml#21376 from ROCm:ci_f8

  Show description for 720d2b8
  

  jax authors

  committed
  720d2b8

  Copy full SHA for 720d2b8

• Merge pull request jax-ml#21434 from jakevdp:take-annotation

  Show description for ebd9c74
  

  jax authors

  committed
  ebd9c74

  Copy full SHA for ebd9c74

• [export] Fix calling under pmap of exported computation with polymorphic shapes

  Show description for acb56a2
  gnecula

  committed
  acb56a2

  Copy full SHA for acb56a2

• Fix rel_entr behavior at boundary value

  sh0416

  committed
  818e7d9

  Copy full SHA for 818e7d9

• [Mosaic GPU] Disable XLA autotuning in matmul_test.py.

  Show description for 72b111a
  

  bchetioui

  authored and

  jax authors

  committed
  72b111a

  Copy full SHA for 72b111a

• [Mosaic] Expand vector.shape_cast no-op detection for expanding/shrinking lane shape casts

  Show description for 97f9a5e
  

  tlongeri

  authored and

  jax authors

  committed
  97f9a5e

  Copy full SHA for 97f9a5e

• [Mosaic] Expand support of vector.broadcast

  Show description for 3fb9acf
  

  tlongeri

  authored and

  jax authors

  committed
  3fb9acf

  Copy full SHA for 3fb9acf

• Merge pull request jax-ml#21448 from gnecula:exp_simpl1

  Show description for e8a1113
  

  jax authors

  committed
  e8a1113

  Copy full SHA for e8a1113

• Update XLA dependency to use revision

  Show description for 149a771
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  149a771

  Copy full SHA for 149a771

• [export] Simplify construction of shardings for the VJP

  Show description for 7dbab16
  gnecula

  committed
  7dbab16

  Copy full SHA for 7dbab16

Commits on May 27, 2024

• Add jax.make_array_from_process_local_data to the docs. Preview: https://jax--21444.org.readthedocs.build/en/21444/_autosummary/jax.make_array_from_process_local_data.html#jax.make_array_from_pro…

  Show description for 08b1cef
  

  yashk2810

  authored and

  jax authors

  committed
  08b1cef

  Copy full SHA for 08b1cef

Commits on May 26, 2024

• Update XLA dependency to use revision

  Show description for 24509f3
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  24509f3

  Copy full SHA for 24509f3

Commits on May 25, 2024

• jnp.take: fix annotation for fill_value

  jakevdp

  committed
  0ff0d7b

  Copy full SHA for 0ff0d7b

• Merge pull request jax-ml#21433 from mattjj:scan-avoid-singleton-dim-in-lowering

  Show description for 5d5ce1c
  

  jax authors

  committed
  5d5ce1c

  Copy full SHA for 5d5ce1c

• avoid singleton dim in scan lowering when unroll==1

  mattjj

  committed
  a24b738

  Copy full SHA for a24b738

• Merge pull request jax-ml#21414 from mattjj:pjit-forwarding-rule

  Show description for 787b7c2
  

  jax authors

  committed
  787b7c2

  Copy full SHA for 787b7c2

• add pjit forwarding rule

  Show description for 0a693fa
  

  mattjj

  and

  froystig

  committed
  0a693fa

  Copy full SHA for 0a693fa

• Update XLA dependency to use revision

  Show description for 9595f14
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  9595f14

  Copy full SHA for 9595f14

Commits on May 24, 2024

• Minor comment correction.

  Show description for c29cff7
  

  Dateng Lin

  authored and

  jax authors

  committed
  c29cff7

  Copy full SHA for c29cff7

• Enable a lot of previously disabled tests now that compute offload works on TPU.

  Show description for c577087
  

  yashk2810

  authored and

  jax authors

  committed
  c577087

  Copy full SHA for c577087

• Add JAX version to TPU_ML_PLATFORM_VERSION environment variable.

  Show description for 93170d9
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  93170d9

  Copy full SHA for 93170d9

• [Pallas][Mosaic] Add lowering rules for PRNG ops.

  Show description for 683ca2c
  

  justinjfu

  authored and

  jax authors

  committed
  683ca2c

  Copy full SHA for 683ca2c

• [xla:python] Add support for registering custom call targets for all XLA execution stages and for XLA FFI traits

  Show description for d5c7ccc
  

  ezhulenev

  authored and

  jax authors

  committed
  d5c7ccc

  Copy full SHA for d5c7ccc

• Update cuDNN to version 9.1.1 in JAX

  Show description for 15a1985
  

  beckerhe

  authored and

  jax authors

  committed
  15a1985

  Copy full SHA for 15a1985

• Merge pull request jax-ml#21262 from vfdev-5:depr-change-ddof-to-correction-21088

  Show description for bab7f40
  

  jax authors

  committed
  bab7f40

  Copy full SHA for bab7f40

• Added correction arg in jnp.var and jnp.std

  Show description for 55f8284
  vfdev-5

  committed
  55f8284

  Copy full SHA for 55f8284

• Add the threefry_partitionable config to JaxprEqnContext to allow setting it inside jit.

  Show description for 972fd66
  

  yashk2810

  authored and

  jax authors

  committed
  972fd66

  Copy full SHA for 972fd66

Pagination

Previous

Next