jax-ml/jax
View on GitHub ↗ ★ 36,567 Python
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
82 mentions on Hacker News by 61 people: 11 stories, 71 comments · first 2018-11-21 · latest 2026-04-05
20082026
-
what do you think of something like https://github.com/google/jax that i guess claims to be a jack of all trades, in the sense that it'll transform programs (through XLA???) to various architectures?
comment by throwlaplace · 2020-04-17 · in A Taste of GPU Compute [video] -
…on different things. We haven't tried combining them yet, but we think it would be fun to explore ( https://github.com/google/jax/issues/1870 ). For example, you could use Numba to hand write a numerical kernel that then participates in a…
comment by emu · 2020-04-08 · in Jax – Composable transformations of Python and NumPy programs -
…so you get those automatically under jax.jit. There's also the jax.checkpoint decorator ( https://github.com/google/jax/pull/1749 ) which lets you control reverse-mode checkpointing yourself; you can use it recursively to implement…
comment by mattjjatgoogle · 2020-03-16 · in Using JAX, numpy, and optimization techniques to improve separable image filters -
…to discover that Jax can't use a TPU's CPU, and that there are no plans to add this to Jax. https://github.com/google/jax/issues/2108#issuecomment-58154... A TPU's CPU is the only reason that TPUs are able to get such high performance on…
comment by sillysaurusx · 2020-02-17 · in Google/Trax – Understand and explore advanced deep learning -
…programs. Design patterns, rules of thumb, and standard practices have yet to be defined. [1] https://github.com/google/jax [2] https://github.com/HIPS/autograd
comment by aaron-santos · 2020-02-16 · in Demystifying Differentiable Programming (2018) -
Do these help? https://github.com/google/jax/tree/master/cloud_tpu_colabs The ODE solver might be close to what you want. I use Tensorflow 1.15. The world has been steadily pushing for Tensorflow 2.0 or Jax, but I like the simplicity of…
comment by sillysaurusx · 2020-01-07 · in AMD’s 64-Core Threadripper 3990X, only $3990 Coming February 7th -
…functions). For example, there must be at least 4 different neural network libraries in JAX. (e.g., https://github.com/google/jax/blob/master/jax/experimental/s... , https://github.com/google/trax ).
comment by nihilest · 2019-12-08 · in Tensorflow User Experience -
…flow interesting since it tries to be more pythonic while also being possible to jit-compile. ( https://github.com/google/jax/blob/master/jax/experimental/l... )
comment by schsam · 2019-12-07 · in Tensorflow User Experience -
I really like the UX of JAX, https://github.com/google/jax The API is numpy plus like four functions. That's the beginning and end. It does JIT compilation under the hood, so can run quite fast.
comment by jlebar · 2019-12-07 · in Tensorflow User Experience -
The practical counterpoint to Dex is Jax, by many of the same authors: https://github.com/google/jax I’m optimistic that we’ll eventually be able to figure out how to build practical tools for deep learning (and scientific computing for…
comment by shoyer · 2019-10-26 · in Dex: A research language for array processing in the Haskell/ML family -
It has stax, which is a minimal example of how to build a high level library: https://github.com/google/jax/blob/master/jax/experimental/s... It seems that the JAX developers are focusing their time on making the core framework better and…
comment by ibab · 2019-10-13 · in Deep Learning with PyTorch: A 60 Minute Blitz [video] -
I really like JAX as well: https://github.com/google/jax . It's younger than PyTorch and TF, but feels cleaner and more expressive. It has a very nice autodiff implementation (based on https://github.com/HIPS/autograd ) and performance is…
comment by ibab · 2019-10-13 · in Deep Learning with PyTorch: A 60 Minute Blitz [video] -
…to transition from TF1 to TF2, consider doing the TF1 to pytorch transition instead. Or consider https://github.com/google/jax !
comment by jlebar · 2019-09-30 · in Tensorflow 2.0 -
For what it's worth, Jax (github.com/google/jax) now lets you use XLA to compile Numpy code. It'd be cool to see how that would stack up in a modern comparison.
comment by davmre · 2019-09-22 · in Weld: Accelerating numpy, scikit and pandas as much as 100x with Rust and LLVM -
How does this compare to Jax? https://github.com/google/jax Why do this in the language instead of a library?
comment by hsaliak · 2019-09-06 · in Differentiable Programming Mega-Proposal -
…case, you should be able to get really good performance without sacrificing readability. Also Jax - https://github.com/google/jax
comment by hsaliak · 2019-08-18 · in Python vs. Rust for Neural Networks -
…capabilities. There is recent work to put numpy on the same footing for deep learning. Check https://github.com/google/jax
comment by jimmy_dean · 2019-08-18 · in Python vs. Rust for Neural Networks -
JAX (github.com/google/jax), which is being developed by many of the authors of Autograd, is a probably a better comparison. At the cost of requiring you to rewrite control flow in a functional way, it eliminates Python's overhead by…
comment by shoyer · 2019-07-19 · in A Differentiable Programming System to Bridge ML and Scientific Computing -
…that there are only 2 people working on the project, if that. The name of the project is JAX: https://github.com/google/jax . It's arguably a lower level framework on top of which something like PyTorch could be built.
comment by m0zg · 2019-01-11 · in NLP Architect by Intel AI Lab -
Check out github.com/google/jax, it’s NumPy on the GPU with automatic differentiation, JIT and autobatching.
comment by alexbw · 2018-12-08 · in PyTorch 1.0 is out -
Those code examples look very clean: https://github.com/google/jax/tree/master/examples
comment by nabla9 · 2018-11-21 · in Compiling machine learning programs via high-level tracing [pdf]