12 Commits

Author	SHA1	Message	Date
jax authors	a207fe9b77	Export KeyPath and related types to jax.tree_util ... These types lie on the APIs in `jax.tree_util`, so it makes sense to export them. PiperOrigin-RevId: 657987755	2024-07-31 06:41:33 -07:00
Jake VanderPlas	d82b66f77f	Document jax.tree.* directly	2024-06-12 14:01:27 -07:00
Jake VanderPlas	1327143d46	Better documentation for jax.tree_util	2024-05-20 19:56:47 -07:00
Jake VanderPlas	defb53f446	Document jax.tree_util.register_dataclass	2024-05-16 18:34:54 -07:00
Tom Hennigan	ed073aa6c9	Add jax.tree_util.tree_leaves_with_path(tree). ... PiperOrigin-RevId: 539609052	2023-06-12 04:13:37 -07:00
ivyzheng	6bf1cbc667	Add key path related guide & code to the documentation.	2023-05-31 20:15:56 -07:00
Cristian Garcia	1e7f9d9e78	add key_path functions to tree_util api reference	2023-03-29 15:17:13 +00:00
Jake VanderPlas	81e627d5bd	DOC: make API doc titles more uniform	2023-01-18 10:59:42 -08:00
Roy Frostig	b8f9dd6269	unify tree_map and tree_multimap	2021-04-28 19:59:31 -07:00
Jake VanderPlas	ffd185612c	DOC: fix tree_util.rst	2021-04-08 15:59:55 -07:00
George Necula	4e89d43a75	Added JAX pytrees notebook ... Also added docstrings to the tree_util module.	2019-11-24 20:29:07 +01:00
Peter Hawkins	8c2c1a2e71	Expose jax.tree_util in the JAX docs.	2019-05-14 21:00:27 -04:00