[D] Current State of JAX vs Pytorch?

• equinox

  31 1,809 9.2 Python

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

There's a good chance that you're interested in building neural networks. In this case be aware that JAX is roughly equivalent to just the torch namespace, and that you can/should choose from various external libraries for building neural networks. The two most popular are Flax and Haiku. Personally I use Equinox which is designed to be a lot more powerful, easier to use, more general etc. (Disclaimer: I am the author of Equinox -- it's something I wrote when I found that Flax/Haiku simply weren't suitable for my use cases.)

• dm-haiku

  10 2,806 7.8 Python

  JAX-based neural network library

  

Just going to add that you should check out haiku if you are considering JAX: https://github.com/deepmind/dm-haiku

• InfluxDB

  www.influxdata.com sponsored

  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.

  

• functorch

  11 1,372 0.6 Jupyter Notebook

  functorch is JAX-like composable function transforms for PyTorch.

  

Fwiw, composable vmap and stuff like that have also been implemented in PyTorch now - see functorch :) https://github.com/pytorch/functorch

Suggest a related project