Peter Hawkins
3c3f01e6d3
Address review comments.
2019-07-30 10:15:37 -04:00
Peter Hawkins
510a9167c5
Add C++ implementation of pytree logic.
...
Move jaxlib version test into jax/lib/__init__.py. Make jax/lib mirror the structure of jaxlib; e.g., xla_client is now available as jax.lib.xla_client.
2019-07-29 15:06:05 -04:00
Peter Hawkins
476dc3db64
Python changes in preparation for adding a C++ implementation of the PyTree utilities.
2019-07-29 10:57:27 -04:00
Matthew Johnson
0546c94992
speed up pmap axis-size getting
...
Co-authored-by: Peter Hawkins <phawkins@google.com>
2019-07-25 12:41:31 -07:00
Peter Hawkins
08013954a4
Use fastcache for LRU caches in JAX.
...
fastcache is both a faster cache implementation and is also thread-safe.
2019-07-22 17:24:10 -04:00
Peter Hawkins
eefd551767
Support older Numpy versions that don't have np.quantile.
...
Fix typo.
2019-07-20 08:44:04 +01:00
Stephan Hoyer
0d6b3f2b18
Define partial to pytree in terms of tuple and dict
2019-07-09 12:05:59 -07:00
Stephan Hoyer
a45fc83eef
Add a functools.partial subclass in tree_util
2019-07-09 11:45:30 -07:00
Matthew Johnson
fadd18b36c
namedtuple subclass transparency ( fixes #806 )
2019-06-03 07:22:32 -07:00
Matthew Johnson
88f691f896
make namedtuples transparent (act as pytree nodes)
2019-05-20 10:08:33 -07:00
Matthew Johnson
642d2dc802
revies optimizers api, fix misc bugs
...
* add more optimizers numerical tests
* update examples and readme with new optimziers api
* add device_values parameter to xla_call
* change optimizers.py to flatten trees and subtrees
* remove tree_map2, tree_multimap2, tree_mimomap, tree_prefixmap
* add optimizer tests: DeviceTuples and error msgs
* make the device_values arg to jit private
2019-05-03 12:44:52 -07:00
Matthew Johnson
87a150e567
add a tree_util.py module-level docstring
2019-05-02 08:02:01 -07:00
Matthew Johnson
50f52070fc
comment-out new pytree check with a TODO
2019-03-22 07:38:40 -07:00
Matthew Johnson
f94c20633a
back to original tree_mimomap optimizers solution
2019-03-22 07:12:25 -07:00
Matthew Johnson
6fdf64d7ee
add back an OptState that is a proper pytree
2019-03-21 19:23:00 -07:00
Matthew Johnson
04cfa11ebe
almost done with a prefix_multimap solution
2019-03-21 19:08:19 -07:00
Matthew Johnson
a20e8982fa
completed scan (PAIR=hawkinsp@)
2019-03-02 21:27:52 -08:00
Matthew Johnson
08dc6994f5
partial progress
2019-02-01 17:05:49 -08:00
Matthew Johnson
0f7c7c4eab
generalize jacfwd and jacrev to handle pytrees
2019-01-06 12:49:41 -08:00
Matthew Johnson
ad4322c5da
playing around with flattening functions
2019-01-06 12:49:35 -08:00
Matthew Johnson
ab53373665
source sync
...
PiperOrigin-RevId: 222470141
2018-11-21 20:32:33 -08:00
Peter Hawkins
5e60639bc5
source sync
...
PiperOrigin-RevId: 222452709
2018-11-21 20:22:54 -08:00
Peter Hawkins
e180f08113
source sync
...
PiperOrigin-RevId: 222451919
2018-11-21 20:22:51 -08:00
Matthew Johnson
a30e858e59
populating source tree
2018-11-17 18:03:33 -08:00
