Commits

Branch selector

 c980dc4

User selector

All users

Datepicker

All time

Commit History

Commits on Apr 19, 2024

• Update outdated comment Pmap_Cookbook

  Show description for c980dc4
  sagelywizard

  authored
  c980dc4

  Copy full SHA for c980dc4

• jax.debug.callback now requires a Callable[..., None]

  Show description for 32922f6
  

  superbobry

  authored and

  jax authors

  committed
  32922f6

  Copy full SHA for 32922f6

• [PJRT C API] Add stream extension to support DLPack and implement this extension in CUDA plugin.

  Show description for b2375fa
  

  Jieying Luo

  authored and

  jax authors

  committed
  b2375fa

  Copy full SHA for b2375fa

• [jax2tf] Adjust tolerance for asinh test.

  Show description for cea36a0
  

  gnecula

  authored and

  jax authors

  committed
  cea36a0

  Copy full SHA for cea36a0

• Merge pull request jax-ml#20825 from jakevdp:gammasgn

  Show description for c7517b8
  

  jax authors

  committed
  c7517b8

  Copy full SHA for c7517b8

• Finalize deprecation of zero-dimensional inputs to jnp.nonzero

  Show description for 41fa67c
  

  Jake VanderPlas

  authored and

  jax authors

  committed
  41fa67c

  Copy full SHA for 41fa67c

• Cache the _check_sharding check in device_put. If aval and sharding are the same, no need to check multiple times

  Show description for 837f0bb
  

  yashk2810

  authored and

  jax authors

  committed
  837f0bb

  Copy full SHA for 837f0bb

• Update XLA dependency to use revision

  Show description for 8fec8a6
  

  jax authors

  committed
  8fec8a6

  Copy full SHA for 8fec8a6

Commits on Apr 18, 2024

• Add jax.scipy.special.gammasgn

  jakevdp

  committed
  568db10

  Copy full SHA for 568db10

• Merge pull request jax-ml#20823 from sagelywizard:patch-1

  Show description for 4d6c4d6
  

  jax authors

  committed
  4d6c4d6

  Copy full SHA for 4d6c4d6

• Remove outdated section in cloud_tpu_colabs README

  Show description for afc8778
  sagelywizard

  authored
  afc8778

  Copy full SHA for afc8778

• [Jax/Triton] Skip benchmarking while autotuning for configs that cannot be launched.

  Show description for 90e9e47
  

  marvin-kim

  authored and

  jax authors

  committed
  90e9e47

  Copy full SHA for 90e9e47

• Merge pull request jax-ml#20816 from superbobry:int4

  Show description for de7d3b6
  

  jax authors

  committed
  de7d3b6

  Copy full SHA for de7d3b6

• Make core.Token a non-trivial class which wraps a jax.Array. Currently, we use a singleton and empty core.token object everywhere. After the change, tokens could be created and threaded in an…

  Show description for c2d4373
  

  yueshengys

  authored and

  jax authors

  committed
  c2d4373

  Copy full SHA for c2d4373

• [Pallas TPU] Generalize while_loop lowering in Pallas -> Mosaic.

  Show description for 9c9e805
  

  jax authors

  committed
  9c9e805

  Copy full SHA for 9c9e805

• Merge pull request jax-ml#20537 from jakevdp:update-tpu-readme

  Show description for 6ca69f3
  

  jax authors

  committed
  6ca69f3

  Copy full SHA for 6ca69f3

• Skip test_spmd_preserves_input_sharding_vmap_grad unless xla_extension_version >= 258

  Show description for aaac1d8
  

  yashk2810

  authored and

  jax authors

  committed
  aaac1d8

  Copy full SHA for aaac1d8

• Fix bug in rank-deficient fix-up code: Do not zero out the corresponding column of u_out if a diagonal entry of r is exactly zero.

  Show description for 51763d8
  

  jax authors

  committed
  51763d8

  Copy full SHA for 51763d8

• Added int4 and uint4 to dtype-specific tests

  Show description for a13efc2
  superbobry

  committed
  a13efc2

  Copy full SHA for a13efc2

• Document the fact that jax.clear_caches() doesn't affect the persistent cache.

  Show description for bb8cf34
  

  jax authors

  committed
  bb8cf34

  Copy full SHA for bb8cf34

• Increase sharding to avoid timeouts

  Show description for fa66e73
  

  apaszke

  authored and

  jax authors

  committed
  fa66e73

  Copy full SHA for fa66e73

• Initial commit for Mosaic GPU

  Show description for 8e3f5b1
  

  apaszke

  authored and

  jax authors

  committed
  8e3f5b1

  Copy full SHA for 8e3f5b1

• Update XLA dependency to use revision

  Show description for c4dea62
  

  jax authors

  committed
  c4dea62

  Copy full SHA for c4dea62

• Merge pull request jax-ml#20806 from jakevdp:tweak-docs

  Show description for 2a65cd3
  

  jax authors

  committed
  2a65cd3

  Copy full SHA for 2a65cd3

Commits on Apr 17, 2024

• DOC: one last readthrough of the new 101 tutorials

  jakevdp

  committed
  48e8457

  Copy full SHA for 48e8457

• Merge pull request jax-ml#20768 from chaserileyroberts:chase/eval_jaxpr_in_extend

  Show description for 1a8aae0
  

  jax authors

  committed
  1a8aae0

  Copy full SHA for 1a8aae0

• Add more imports to jax extend

  chaserileyroberts

  committed
  74c2e25

  Copy full SHA for 74c2e25

• [XLA:Mosaic] Generalize (8,128) -> (8 * packing,128) retiling for packed type.

  Show description for d44b16c
  

  bythew3i

  authored and

  jax authors

  committed
  d44b16c

  Copy full SHA for d44b16c

• Remove the spmd_mode check from OSS JAX since enhanced barrier is switched on for OSS JAX

  Show description for 7cb0e60
  

  yashk2810

  authored and

  jax authors

  committed
  7cb0e60

  Copy full SHA for 7cb0e60

• Merge pull request jax-ml#20802 from jakevdp:dup-thinking-in-jax

  Show description for 1c8534e
  

  jax authors

  committed
  1c8534e

  Copy full SHA for 1c8534e

• Merge pull request jax-ml#20800 from jakevdp:pin-sphinx

  Show description for 73bd54f
  

  jax authors

  committed
  73bd54f

  Copy full SHA for 73bd54f

• DOC: remove copy of thinking-in-jax from new tutorial flow

  jakevdp

  committed
  edfdc36

  Copy full SHA for edfdc36

• DOC: pin sphinx to >=7.3.2

  jakevdp

  committed
  29b15e2

  Copy full SHA for 29b15e2

• Fix Pallas' while_loop lowering to properly account for the loop length

  Show description for acde885
  

  apaszke

  authored and

  jax authors

  committed
  acde885

  Copy full SHA for acde885

• Reverts f187399

  Show description for 6963c77
  

  jax authors

  committed
  6963c77

  Copy full SHA for 6963c77

Pagination

Previous

Next