Pytrees

• tree

  5 906 4.7 Python

  tree is a library for working with nested data structures (by google-deepmind)

  

There is also the standalone library "tree" from DeepMind: https://github.com/deepmind/tree

It provides similar functionality but is standalone and does not depend on JAX, TF or anything else.

• equinox

  31 1,819 9.2 Python

  Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

You're thinking of `jax.closure_convert`. :)

(Although technically that works by tracing and extracting all constants from the jaxpr, rather than introspecting the function's closure cells -- it sounds like your trick is the latter.)

When you discuss dynamic allocation, I'm guessing you're mainly referring to not being able to backprop through `jax.lax.while_loop`. If so, you might find `equinox.internal.while_loop` interesting, which is an unbounded while loop that you can backprop through! The secret sauce is to use a treeverse-style checkpointing scheme.

https://github.com/patrick-kidger/equinox/blob/f95a8ba13fb35...

• InfluxDB

  www.influxdata.com featured

  Power Real-Time Data Analytics at Scale. Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.

  

Suggest a related project