Commits

Branch selector

 main

User selector

mattjj

Datepicker

All time

Commit History

Commits on Sep 9, 2022

• [dynamic-shapes] add basic vmap-of-indexing support

  Show description for 5882650
  mattjj

  committed
  5882650

  Copy full SHA for 5882650

Commits on Sep 8, 2022

• add _device attribute to PRNGKeyArray so that computation follows key placement

  Show description for 47b2dfe
  mattjj

  committed
  47b2dfe

  Copy full SHA for 47b2dfe

Commits on Sep 6, 2022

• fix bugs, infeed/outfeed must be considered effectful

  Show description for 3c811b1
  

  3 peoplecommitted
  3c811b1

  Copy full SHA for 3c811b1

Commits on Aug 29, 2022

• Add batching rules for state primitives and for_loop

  Show description for bbb8048
  

  mattjj

  and

  sharadmv

  committed
  bbb8048

  Copy full SHA for bbb8048

Commits on Aug 24, 2022

• in ad_checkpoint WrapHashably.__eq__, check _both_ are hashable

  Show description for 88a212e
  mattjj

  committed
  88a212e

  Copy full SHA for 88a212e

Commits on Aug 19, 2022

• make slicing compilation a little faster

  Show description for b3b4ffb
  mattjj

  committed
  b3b4ffb

  Copy full SHA for b3b4ffb

Commits on Aug 18, 2022

• rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

  Show description for b7426b5
  

  mattjj

  authored and

  jax authors

  committed
  b7426b5

  Copy full SHA for b7426b5

• clean up some dead code, from bad merge

  mattjj

  committed
  d2647cf

  Copy full SHA for d2647cf

• roll-forward deletion of custom_jvp_call_jaxpr_p

  Show description for 3a20de1
  

  mattjj

  authored and

  jax authors

  committed
  3a20de1

  Copy full SHA for 3a20de1

• remove custom_jvp_call_jaxpr_p and its rules

  Show description for 887b7ce
  mattjj

  committed
  887b7ce

  Copy full SHA for 887b7ce

Commits on Aug 17, 2022

• delete old remat implementation

  Show description for d19e34f
  mattjj

  committed
  d19e34f

  Copy full SHA for d19e34f

Commits on Aug 15, 2022

• un-skip polar/qdwh decomp tests skipped on gpu in ad6ce74

  Show description for 68e3f58
  mattjj

  committed
  68e3f58

  Copy full SHA for 68e3f58

• Working multihost eager pmap

  Show description for a7f760d
  

  3 peoplecommitted
  a7f760d

  Copy full SHA for a7f760d

• Initial implementation of eager pmap

  Show description for 5310515
  

  mattjj

  and

  sharadmv

  committed
  5310515

  Copy full SHA for 5310515

Commits on Aug 13, 2022

• simplify slicing jaxprs a little

  Show description for 42dd7ca
  

  mattjj

  and

  sharadmv

  committed
  42dd7ca

  Copy full SHA for 42dd7ca

Commits on Aug 10, 2022

• set new jax.remat / jax.checkpoint to be on-by-default

  mattjj

  committed
  be6f6bf

  Copy full SHA for be6f6bf

• fix type annotation on remat

  mattjj

  committed
  d76754e

  Copy full SHA for d76754e

Commits on Aug 9, 2022

• fix traceback exclusion on new checkpoint

  mattjj

  committed
  580e7f3

  Copy full SHA for 580e7f3

• fix diffrax with new remat, do for cond what jax-ml#11773 did for while_loop

  mattjj

  committed
  666f4f8

  Copy full SHA for 666f4f8

Commits on Aug 6, 2022

• Rolling forward jax-ml#11768 after test failures caused roll-back (from use of np.empty).

  Show description for 81b6263
  

  mattjj

  authored and

  jax authors

  committed
  81b6263

  Copy full SHA for 81b6263

• add api_boundary, fix flax failure with new-remat-by-default

  mattjj

  committed
  e2c1df6

  Copy full SHA for e2c1df6

• handle test failure in diffrax under new remat

  Show description for 9f826d1
  mattjj

  committed
  9f826d1

  Copy full SHA for 9f826d1

• prototype unfettered element types in jaxpr arrays

  Show description for 348da51
  

  mattjj

  and

  froystig

  committed
  348da51

  Copy full SHA for 348da51

Commits on Aug 5, 2022

• add trivial convert_element_type padding rule

  mattjj

  committed
  f159cdf

  Copy full SHA for f159cdf

• small tweaks for bint ad

  mattjj

  committed
  fbf6aa2

  Copy full SHA for fbf6aa2

Commits on Aug 4, 2022

• prepare to switch to new remat

  Show description for 33207e2
  mattjj

  committed
  33207e2

  Copy full SHA for 33207e2

• Update jax/_src/ad_checkpoint.py

  Show description for 1a01e56
  

  mattjj

  and

  froystig

  authored
  1a01e56

  Copy full SHA for 1a01e56

• Update jax/_src/ad_checkpoint.py

  Show description for 9d25fa4
  

  mattjj

  and

  froystig

  authored
  9d25fa4

  Copy full SHA for 9d25fa4

• prepare to switch to new remat

  Show description for e3a92d5
  mattjj

  committed
  e3a92d5

  Copy full SHA for e3a92d5

• in sharp bits notebook, add back iaml blog post link

  mattjj

  committed
  7ed4fa5

  Copy full SHA for 7ed4fa5

Commits on Jul 30, 2022

• fix ad_checkpoint.checkpoint caching issue

  Show description for cbcfe95
  mattjj

  committed
  cbcfe95

  Copy full SHA for cbcfe95

Commits on Jul 29, 2022

• add test for custom_linear_solve + new remat

  mattjj

  committed
  aa043a6

  Copy full SHA for aa043a6

• add custom-policy partial eval and dce rules for pmap

  Show description for e0c1e6c
  mattjj

  committed
  e0c1e6c

  Copy full SHA for e0c1e6c

• add while_loop custom-policy partial eval rule

  mattjj

  committed
  7f3aa12

  Copy full SHA for 7f3aa12

Commits on Jul 28, 2022

• add cond dce rule and custom-policy partial eval rule

  mattjj

  committed
  ec9f9c3

  Copy full SHA for ec9f9c3

Pagination

Previous

Next