7 Commits

Author	SHA1	Message	Date
Peter Hawkins	38bffe9a8b	Add a pytreedef.flatten_up_to() method that flattens a PyTree only up to the structure of a PyTreeDef. ... Make the C++ version of tree_multimap accept tree suffixes of the primary tree. Document and test this behavior. Remove unnecessary locking in custom node registry; we hold the GIL already so there's no point to the additional locking.	2019-08-01 12:17:00 -04:00
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	5ceac99d0c	Add newline to the end of jaxlib/BUILD.	2019-04-01 08:23:02 -07:00
Peter Hawkins	5eff830f0e	Move jaxlib version.py into jaxlib, and install it in build/jaxlib as build action. ... Update jaxlib version check to look in jaxlib.version.	2019-04-01 08:21:22 -07:00
Peter Hawkins	20935448b9	Whitespace fix to jaxlib/BUILD.	2018-12-19 14:51:46 -05:00
Peter Hawkins	3c388b98f1	Add support for calling LAPACK primitives from SciPy from JAX linalg.	2018-12-17 16:30:27 -05:00