functorch
equinox
functorch | equinox | |
---|---|---|
11 | 31 | |
1,372 | 1,819 | |
0.4% | - | |
0.0 | 9.2 | |
1 day ago | 15 days ago | |
Jupyter Notebook | Python | |
BSD 3-clause "New" or "Revised" License | Apache License 2.0 |
Stars - the number of stars that a project has on GitHub. Growth - month over month growth in stars.
Activity is a relative number indicating how actively a project is being developed. Recent commits have higher weight than older ones.
For example, an activity of 9.0 indicates that a project is amongst the top 10% of the most actively developed projects that we are tracking.
functorch
-
What is the most efficient approach to ensemble a pytorch actor-critic model?
I would suggest checking https://pytorch.org/functorch/ and https://github.com/metaopt/torchopt for efficient inference and training with ensembles (e.g., t be independent actors in a multi-agent setting or multiple critics).
-
[P] Multidimensional array batch indexing for pytorch and numpy
There were some bugs still with advanced indexing in an older release of functorch, I believe they should be fixed now though: https://github.com/pytorch/functorch/pull/862
- Functorch: Jax-like composable function transforms for PyTorch
-
Jax vs. Julia (Vs PyTorch)
Tangentially related but there is an effort to get some of the features of JAX into PyTorch: https://pytorch.org/functorch/
-
[D] Current State of JAX vs Pytorch?
Fwiw, composable vmap and stuff like that have also been implemented in PyTorch now - see functorch :) https://github.com/pytorch/functorch
-
[D] Ideal deep learning library
Fwiw, it’s not like Pytorch’s design prevents function transformations from being implemented. See functorch for an example of grad/vmap function transforms: https://github.com/pytorch/functorch
-
[P] Made Some Pytorch Modules For Agent Systems
You may find vmap from functorch to be quite useful: https://github.com/pytorch/functorch
-
[D] Are you using PyTorch or TensorFlow going into 2022?
If you're interested in function transformations in PyTorch, try out functorch :) https://github.com/pytorch/functorch
- PyTorch: Where we are headed and why it looks a lot like Julia (but not exactly)
-
Show HN: How does Jax allocate memory on a TPU? An interactive C++ walkthrough
The pytorch programming model is just really hard to adapt to an XLA-like compiler. Imperative python code doesn't translate to an ML graph compiler particularly well; Jax's API is functional, so it's easier to translate to the XLA API. By contrast, torch/xla uses "lazy tensors" that record the computation graph and compile when needed. The trouble is, if the compute graph changes from run to run, you end up recompiling a lot.
I guess in Jax you'd just only apply `jax.jit` to the parts where the compute graph is static? I'd be curious to see examples of how this works in practice. Fwiw, there's an offshoot of pytorch that is aiming to provide this sort of API (see https://github.com/pytorch/functorch and look at eager_compilation.py).
(Disclaimer: I worked on this until quite recently.)
equinox
-
Ask HN: What side projects landed you a job?
I wrote a JAX-based neural network library (Equinox [1]) and numerical differential equation solving library (Diffrax [2]).
At the time I was just exploring some new research ideas in numerics -- and frankly, procrastinating from writing up my PhD thesis!
But then one of the teams at Google starting using them, so they offered me a job to keep developing them for their needs. Plus I'd get to work in biotech, which was a big interest of mine. This was a clear dream job offer, so I accepted.
Since then both have grown steadily in popularity (~2.6k GitHub stars) and now see pretty widespread use! I've since started writing several other JAX libraries and we now have a bit of an ecosystem going.
[1] https://github.com/patrick-kidger/equinox
-
[P] Optimistix, nonlinear optimisation in JAX+Equinox!
The elevator pitch is Optimistix is really fast, especially to compile. It plays nicely with Optax for first-order gradient-based methods, and takes a lot of design inspiration from Equinox, representing the state of all the solvers as standard JAX PyTrees.
-
JAX – NumPy on the CPU, GPU, and TPU, with great automatic differentiation
If you like PyTorch then you might like Equinox, by the way. (https://github.com/patrick-kidger/equinox ; 1.4k GitHub stars now!)
- Equinox: Elegant easy-to-use neural networks in Jax
- Show HN: Equinox (1.3k stars), a JAX library for neural networks and sciML
-
Pytrees
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...
-
Writing Python like it’s Rust
I'm a big fan of using ABCs to declare interfaces -- so much so that I have an improved abc.ABCMeta that also handles abstract instance variables and abstract class variables: https://github.com/patrick-kidger/equinox/blob/main/equinox/_better_abstract.py
-
[D] JAX vs PyTorch in 2023
For the daily research, I use Equinox (https://github.com/patrick-kidger/equinox) as a DL librarry in JAX.
- [Machinelearning] [D] État actuel de JAX vs Pytorch?
-
Training Deep Networks with Data Parallelism in Jax
It sounds like you're concerned about how downstream libraries tend to wrap JAX transformations to handle their own thing? (E.g. `haiku.grad`.)
If so, then allow me to make my usual advert here for Equinox:
https://github.com/patrick-kidger/equinox
This actually works with JAX's native transformations. (There's no `equinox.vmap` for example.)
On higher-order functions more generally, Equinox offers a way to control these quite carefully, by making ubiquitous use of callables that are also pytrees. E.g. a neural network is both a callable in that it has a forward pass, and a pytree in that it records its parameters in its tree structure.
What are some alternatives?
jax - Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
flax - Flax is a neural network library for JAX that is designed for flexibility.
nn - 🧑🏫 60 Implementations/tutorials of deep learning papers with side-by-side notes 📝; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), 🎮 reinforcement learning (ppo, dqn), capsnet, distillation, ... 🧠
dm-haiku - JAX-based neural network library
onnx-simplifier - Simplify your onnx model
torchtyping - Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.
torch2trt - An easy to use PyTorch to TensorRT converter
treex - A Pytree Module system for Deep Learning in JAX
BinaryBuilder.jl - Binary Dependency Builder for Julia
extending-jax - Extending JAX with custom C++ and CUDA code
py2many - Transpiler of Python to many other languages
diffrax - Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/