dm-haiku
functorch
dm-haiku | functorch | |
---|---|---|
10 | 11 | |
2,806 | 1,372 | |
0.9% | 0.4% | |
7.8 | 0.0 | |
27 days ago | 3 days ago | |
Python | Jupyter Notebook | |
Apache License 2.0 | BSD 3-clause "New" or "Revised" License |
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.
dm-haiku
-
Maxtext: A simple, performant and scalable Jax LLM
Is t5x an encoder/decoder architecture?
Some more general options.
The Flax ecosystem
https://github.com/google/flax?tab=readme-ov-file
or dm-haiku
https://github.com/google-deepmind/dm-haiku
were some of the best developed communities in the Jax AI field
Perhaps the “trax” repo? https://github.com/google/trax
Some HF examples https://github.com/huggingface/transformers/tree/main/exampl...
Sadly it seems much of the work is proprietary these days, but one example could be Grok-1, if you customize the details. https://github.com/xai-org/grok-1/blob/main/run.py
-
Help with installing python packages.
I am fresh to nix os especially when it comes to using python on it how do I install packages withought using pip I need to install numpy~=1.19.5 transformers~=4.8.2 tqdm~=4.45.0 setuptools~=51.3.3 wandb>=0.11.2 einops~=0.3.0 requests~=2.25.1 fabric~=2.6.0 optax==0.0.6 git+https://github.com/deepmind/dm-haiku git+https://github.com/EleutherAI/lm-evaluation-harness/ ray[default]==1.4.1 jax~=0.2.12 Flask~=1.1.2 cloudpickle~=1.3.0 tensorflow-cpu~=2.5.0 google-cloud-storage~=1.36.2 smart_open[gcs] func_timeout ftfy fastapi uvicorn lm_dataformat which I can just do pip -r thetxtfile but idk how to do this in nix os also I would be using python3.7 so far this is what I have come up with but I know its wrong { pkgs ? import {} }: let packages = python-packages: with python-packages; [ mesh-transformer-jax/ jax==0.2.12 numpy~=1.19.5 transformers~=4.8.2 tqdm~=4.45.0 setuptools~=51.3.3 wandb>=0.11.2 einops~=0.3.0 requests~=2.25.1 fabric~=2.6.0 optax==0.0.6 #the other packages ]; pkgs.mkShell { nativeBuildInputs = [ pkgs.buildPackages.python37 ]; }
-
[D] Should We Be Using JAX in 2022?
What's your favorite Deep Learning API for JAX - Flax, Haiku, Elegy, something else?
-
[D] Current State of JAX vs Pytorch?
Just going to add that you should check out haiku if you are considering JAX: https://github.com/deepmind/dm-haiku
-
PyTorch vs. TensorFlow in 2022
As a researcher in RL & ML in a big industry lab, I would say most of my colleagues are moving to JAX 0https://github.com/google/jax], which this article kind of ignores. JAX is XLA-accelerated NumPy, it's cool beyond just machine learning, but only provides low-level linear algebra abstractions. However you can put something like Haiku [https://github.com/deepmind/dm-haiku] or Flax [https://github.com/google/flax] on top of it and get what the cool kids are using :)
-
[D] JAX learning resources?
- https://github.com/deepmind/dm-haiku/tree/main/examples
- Why would I want to develop yet another deep learning framework?
- Help with installing python packages
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.)
What are some alternatives?
flax - Flax is a neural network library for JAX that is designed for flexibility.
jax - Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
jax-resnet - Implementations and checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX (Flax).
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, ... 🧠
trax - Trax — Deep Learning with Clear Code and Speed
onnx-simplifier - Simplify your onnx model
equinox - Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
torch2trt - An easy to use PyTorch to TensorRT converter
elegy - A High Level API for Deep Learning in JAX
BinaryBuilder.jl - Binary Dependency Builder for Julia
py2many - Transpiler of Python to many other languages