jax authors
|
398aaaacc7
|
Add support for Ellipsis as an index for stateful operations.
PiperOrigin-RevId: 497466879
|
2022-12-23 22:46:50 -08:00 |
|
Jake VanderPlas
|
4a6bbde409
|
Move jax.linear_util to jax._src.linear_util
|
2022-12-20 14:49:27 -08:00 |
|
Jake VanderPlas
|
4389216d0c
|
Remove typing_extensions dependency
|
2022-12-05 15:42:26 -08:00 |
|
Sharad Vikram
|
e1af93a9ba
|
Enable state effect in cond_p (except in grad and vmap)
PiperOrigin-RevId: 485719926
|
2022-11-02 16:07:58 -07:00 |
|
Jake VanderPlas
|
97b17af5be
|
[typing] add type annotations to the first several lax_numpy functions
|
2022-10-21 11:59:53 -07:00 |
|
Jake VanderPlas
|
7f89fd40a2
|
Cleanup: remove unused imports in private modules
Also improve our flake8 filter rules to avoid ignoring these.
|
2022-10-20 14:37:21 -07:00 |
|
Jake VanderPlas
|
5d15757741
|
[typing] annotate jax._src.util.safe_map
|
2022-10-20 10:15:04 -07:00 |
|
Jake VanderPlas
|
524745f322
|
TMP: annotate util.safe_zip
|
2022-10-19 10:29:53 -07:00 |
|
Sharad Vikram
|
bbf69d10cc
|
Enable partially discharging state effects from jaxprs
|
2022-10-11 16:52:29 -07:00 |
|
Peter Hawkins
|
ba557d5e1b
|
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.
PiperOrigin-RevId: 476167538
|
2022-09-22 12:27:19 -07:00 |
|
Sharad Vikram
|
f26f1e8afc
|
Add support for closing over Refs in nested for loops
|
2022-09-13 13:32:44 -07:00 |
|
Sharad Vikram
|
6967c7ef51
|
Add sound loop invariance detection
|
2022-09-08 10:42:19 -07:00 |
|
Sharad Vikram
|
b6c3b9df19
|
Split State effect into Read/Write/Accum effects and tie them to Ref avals
|
2022-09-08 08:04:13 -07:00 |
|
Sharad Vikram
|
b2a5d2c3bb
|
Add partial_eval_custom rule for for_loop
|
2022-09-06 11:00:26 -07:00 |
|
Matthew Johnson
|
bbb8048d2e
|
Add batching rules for state primitives and for_loop
Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
|
2022-08-29 11:40:09 -07:00 |
|
jax authors
|
5d9dfbb142
|
Merge pull request #11900 from sharadmv:for-loop
PiperOrigin-RevId: 468608666
|
2022-08-18 20:14:17 -07:00 |
|
Sharad Vikram
|
49b7729f6b
|
More tests for transpose
|
2022-08-18 18:06:21 -07:00 |
|
Neil Girdhar
|
ad38a6bb28
|
Fix common typo: Tuple[X] -> Tuple[X, ...]
|
2022-08-16 11:47:22 -04:00 |
|
Sharad Vikram
|
72dbe31172
|
Initial transpose implementation
Co-authored-by: Matthew Johnson <mattjj@google.com>
|
2022-08-15 10:23:04 -07:00 |
|
Sharad Vikram
|
8b7daa8095
|
Refactor state out of for_loop
|
2022-08-01 15:26:55 -07:00 |
|