Commits

Branch selector

 main

User selector

All users

Datepicker

All time

Commit History

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

• 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

• Merge pull request jax-ml#21362 from vfdev-5:fix-array-api-compat-var-nan-21330

  Show description for 9b0d403
  

  jax authors

  committed
  9b0d403

  Copy full SHA for 9b0d403

• Merge pull request jax-ml#21423 from hawkinsp:cumsum

  Show description for 31d4a1e
  

  jax authors

  committed
  31d4a1e

  Copy full SHA for 31d4a1e

• Merge pull request jax-ml#21406 from superbobry:maint-too

  Show description for d910c5c
  

  jax authors

  committed
  d910c5c

  Copy full SHA for d910c5c

• Another attempt to land jax-ml#20445

  Show description for 15b974c
  

  superbobry

  authored and

  jax authors

  committed
  15b974c

  Copy full SHA for 15b974c

• Bump up memory requirements for pmap_test on TPU.

  Show description for 1e01fa7
  

  hawkinsp

  authored and

  jax authors

  committed
  1e01fa7

  Copy full SHA for 1e01fa7

• Add note to release notes about jax-ml#21403.

  Show description for 441ab58
  hawkinsp

  committed
  441ab58

  Copy full SHA for 441ab58

• jnp.var returns nan if N-ddof <= 0

  Show description for 3c201e0
  vfdev-5

  committed
  3c201e0

  Copy full SHA for 3c201e0

• [Mosaic GPU] Remove the testonly test_util import from matmul.py (better be imported on the test side)

  Show description for be7939a
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  be7939a

  Copy full SHA for be7939a

• Merge pull request jax-ml#21404 from superbobry:maint

  Show description for e86c436
  

  jax authors

  committed
  e86c436

  Copy full SHA for e86c436

• Removed mentions of iree from the test suite

  superbobry

  committed
  2473ebf

  Copy full SHA for 2473ebf

• [JAX] Fix multi-process data sharing comments.

  Show description for e0bf783
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  e0bf783

  Copy full SHA for e0bf783

• Update XLA dependency to use revision

  Show description for 9140ea3
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  9140ea3

  Copy full SHA for 9140ea3

Commits on May 23, 2024

• Drop refs to australis in build rules and exporter.

  Show description for c1f5a32
  

  pschuh

  authored and

  jax authors

  committed
  c1f5a32

  Copy full SHA for c1f5a32

• Fix the flaky pax checkpoint test by rolling back the jnp.asarray conversion. Also fix the int4 layout breakage.

  Show description for 3dcfe34
  

  yashk2810

  authored and

  jax authors

  committed
  3dcfe34

  Copy full SHA for 3dcfe34

• Bumped the minimum ml_dtypes version to 0.4.0

  superbobry

  committed
  0a694a1

  Copy full SHA for 0a694a1

• Improve error message when trying to fetch value of non-addressable array.

  Show description for b527e1e
  

  yashk2810

  authored and

  jax authors

  committed
  b527e1e

  Copy full SHA for b527e1e

• [Pallas] Add lower precision for pallas out-of-bounds interpret mode test on GPU.

  Show description for ee79d7d
  

  justinjfu

  authored and

  jax authors

  committed
  ee79d7d

  Copy full SHA for ee79d7d

• Update XLA dependency to use revision

  Show description for 8a3bc5d
  

  beckerhe

  authored and

  jax authors

  committed
  8a3bc5d

  Copy full SHA for 8a3bc5d

• Added support for printing scalar values in Pallas TPU kernels

  Show description for daa81e6
  

  superbobry

  authored and

  jax authors

  committed
  daa81e6

  Copy full SHA for daa81e6

Pagination

Previous

Next