extending-jax VS trax

Compare extending-jax vs trax and see what are their differences.

Scout Monitoring - Free Django app performance insights with Scout Monitoring
Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.
www.scoutapm.com
featured
InfluxDB - 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.
www.influxdata.com
featured
extending-jax trax
2 7
363 8,011
- 0.5%
3.5 3.7
9 months ago 6 days ago
Python Python
MIT License Apache License 2.0
The number of mentions indicates the total number of mentions that we've tracked plus the number of user suggested alternatives.
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.

extending-jax

Posts with mentions or reviews of extending-jax. We have used some of these posts to build our list of alternatives and similar projects. The last one was on 2022-02-15.

trax

Posts with mentions or reviews of trax. We have used some of these posts to build our list of alternatives and similar projects. The last one was on 2024-04-23.
  • Maxtext: A simple, performant and scalable Jax LLM
    10 projects | news.ycombinator.com | 23 Apr 2024
    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

  • Replit's new Code LLM was trained in 1 week
    12 projects | news.ycombinator.com | 3 May 2023
    and the implementation https://github.com/google/trax/blob/master/trax/models/resea... if you are interested.

    Hope you get to look into this!

  • RedPajama: Reproduction of Llama with Friendly License
    4 projects | news.ycombinator.com | 17 Apr 2023
    Thank you for developing the pipeline and amassing considerable compute for gathering and preprocessing this dataset!

    I'm not sure if this is the right place to ask about this, but could you consider training an LLM using a more advanced, sparse transformer architecture (specifically, "Terraformer" from this paper https://arxiv.org/abs/2111.12763 and this codebase https://github.com/google/trax/blob/master/trax/models/resea... by Google Brain and OpenAI)? I understand the pressure to focus on training a straightforward LLaMA replication, but of course you see that it's a legacy dense architecture which limits its inference performance. This new architecture is not just an academic curiosity but is already validated at scale by Google, providing 10x+ inference performance boost on the same hardware.

    Frankly, the community's compute budget - for training and for inference - isn't infinite, and neither is the public's interest in models that do not have advantage (at least in convenience) over closed-source ones; and so we should utilize both those resources as efficiently as possible. It could be a big step forward if you trained at least LLaMA-Terraformer-7B and 13B foundation models on the whole dataset.

  • The founder of Gmail claims that ChatGPT can “kill” Google in two years.
    1 project | /r/Futurology | 31 Jan 2023
    But a couple years later they came out with open source implementations yeah: https://github.com/google/trax/tree/master/trax/models/reformer
  • [D] Paper Explained - Sparse is Enough in Scaling Transformers (aka Terraformer) | Video Walkthrough
    1 project | /r/MachineLearning | 1 Dec 2021
    Code: https://github.com/google/trax/blob/master/trax/examples/Terraformer_from_scratch.ipynb
  • Why would I want to develop yet another deep learning framework?
    4 projects | /r/learnmachinelearning | 16 Sep 2021
  • How to train large models on a normal laptop?
    1 project | /r/LanguageTechnology | 14 Feb 2021
    Training language models is expensive. Train the biggest model you can afford. I assume you've tried the colab from the reformer GitHub: https://github.com/google/trax/tree/master/trax/models/reformer

What are some alternatives?

When comparing extending-jax and trax you can also consider the following projects:

einops - Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

flax - Flax is a neural network library for JAX that is designed for flexibility.

mpi4jax - Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python :zap:

dm-haiku - JAX-based neural network library

thinc - 🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

muzero-general - MuZero

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

ML-Optimizers-JAX - Toy implementations of some popular ML optimizers using Python/JAX

elegy - A High Level API for Deep Learning in JAX

objax

diffrax - Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

numpyro - Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

Scout Monitoring - Free Django app performance insights with Scout Monitoring
Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.
www.scoutapm.com
featured
InfluxDB - 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.
www.influxdata.com
featured