JAX – Augments numpy and Python code with function transformations (2019)

https://colinraffel.com/blog/you-don-t-know-jax.html

by auraham • 3 years ago

65 12 3 years ago

l1n• 3 years ago

Needs a [2019]

pestatije• 3 years ago

JAX is a Python library which augments numpy and Python code with function transformations which make it trivial to perform operations common in machine learning programs.

tpoacher• 3 years ago

why?

Q6T46nT668w6i3m2• 3 years ago

NumPy doesn’t provide automatic differentiation or heterogenous compute.

funks_• 3 years ago

JAX is just numpy (mostly)

Or you could say JAX is a DSL that looks like numpy to trick Python programmers into using a compiled functional programming language. ;)

fnands• 3 years ago

Shhh, don't tell them! ;-)

bouchard• 3 years ago

I've recently started using Taichi (https://taichi-lang.org/) for numerical codes and the fact it doesn't try to trick you into thinking it's numpy is a nice "feature". ;)

thanatropism• 3 years ago

Such an inciting title. I "know" Jax; I also know Jack from Legal.

I'd been using Jax and Numpyro like learning from tutorials pretending it's numpy and then they brought the pros from Dover who knew all kinds of insane advanced techniques. I was hoping for a tutorial on those.

imurray• 3 years ago

While I didn't like the form of the title, I assumed it was referring to other tutorials with similar titles. Indeed there exists "you don't know JS" and "you don't know bash".

fernandotakai• 3 years ago

We'll use a neural network with a single hidden layer with 3 neurons and a hyperbolic tangent nonlinearity, trained with the cross-entropy loss via stochastic gradient descent.

i gotta be super honest, this reads like a sentence written by ChatGPT.

eternalban• 3 years ago

I have be super honest too, as your comment says more about your English skills than whether writing a coherent sentence is beyond the reach of humans. Nothing in that sentence says machine generated. But whoever wrote it thinks clearly and precisely, that much is obvious.

Q6T46nT668w6i3m2• 3 years ago

I love JAX … but JAX, frustratingly, doesn’t obey NumPy promotion semantics! Enforcing double precision is understandable from a performance perspective but it breaks NumPy compatibility and frequently introduces subtle and difficult to debug numerical bugs. Nevertheless, JAX has encouraged so much progress in ML software engineering, e.g., I’ve been cribbing from diffrax and torchopt for my PyTorch work.

JAX – Augments numpy and Python code with function transformations (2019)

Related Stories