DOC: replace old tutorials with new content

[jakevdp]

This is the first part of the restructuring in #18585

The main change here is that the older jax-101 tutorials are replaced by newer, updated discussions, with redirects from the old URLs.

Preview here: https://jax--20819.org.readthedocs.build/en/20819/

Find the main new content listed here: https://jax--20819.org.readthedocs.build/en/20819/tutorials.html, and linked in the side-bar.

I've added the new sphinxext-rediraffe extension to add automatic redirects for the old URLs to the replacement content. Some examples of the new redirects in action:

• https://jax--20819.org.readthedocs.build/en/20819/jax-101/ -> https://jax--20819.org.readthedocs.build/en/20819/tutorials.html
• https://jax--20819.org.readthedocs.build/en/20819/jax-101/06-parallelism.html -> https://jax--20819.org.readthedocs.build/en/20819/sharded-computation.html
• https://jax--20819.org.readthedocs.build/en/20819/notebooks/quickstart.html -> https://jax--20819.org.readthedocs.build/en/20819/quickstart.html

[mattjj]

Not sure how to attach this to the right place, but: there's a line at the top of the quickstart that says "a la NumPy", can we change that to "à la NumPy"? I don't want to offend the French.

[mattjj]

(Apologies again, I'm not sure how to attach this comment to the right place due to how the file move is represented in GitHub.)

In quickstart.md, there's a line that says "the jax.jacobian() transformation can be used to compute gradients of vector-valued functions". It refers to the gradients of vector-valued functions.

IMO we should reserve "gradient" only to apply to scalar-valued functions. That's consistent with Wikipedia, the jax.grad restriction, and my understanding of usual conventions.

How about we rephrase to something like:

"Beyond scalar-valued functions, we can use jax.jacobian to compute the full Jacobian matrix of vector-valued functions."

[mattjj]

In quickstart.md there's a sentence that says

For more advanced autodiff operations, you can use jax..jacrev() for reverse-mode vector-Jacobian products, and jax.jacfwd() for forward-mode Jacobian-vector products

But actually 'vector-Jacobian products' and 'jacobian-vector products' are describing jax.vjp and jax.jvp, not jax.jacrev and jax.jacfwd. Let's change it to:

For more advanced autodiff operations, you can use jax.vjp() for reverse-mode vector-Jacobian products, and jax.jvp() or jax.linearize for forward-mode Jacobian-vector products

But then the text that just follows needs to be revised too, since it refers to jacfwd and jacrev... we could define those inline with vmap as ~1 liners each, or add a sentence introducing jacfwd and jacrev separately. I kind of lean towards defining them inline. WDYT?

[mattjj]

In key-concepts.md, there's a line "JAX arrays are never constructed directly" that wasn't immediately clear to me what it meant. How about we change it to something like "We don't call the jax.Array object constructor directly, like jax.Array(...), but rather ..."?

[mattjj]

In key-concepts.md I noticed some "jax" instead of "JAX", but we can fix those later.

I also noticed that "The magic behind transformations is the notion of a Tracers." should end with the word "Tracer" rather than "Tracers" I think (or the "a" should be deleted).

[mattjj]

Maybe I'll try attaching comments to the deleted side of the moved files... EDIT nope couldn't do it, e.g. nowhere to attach comments to key-concepts.md line-by-line afaict.

[mattjj]

Docs are such a great opportunity to check over the fundamentals.

In key-concepts.md, the Tracing section reminded me that I had a recent experiment trying to improve how Tracers are printed. The Traced<ShapedArray(int32[5])>with<DynamicJaxprTrace(level=1/0)> is pretty opaque to users.

This branch makes some small tweaks, so that when we run this test file:

import jax
import jax.numpy as jnp
jax.config.update('jax_platforms', 'cpu')

cool = []
def f(x):
  cool.append(repr(x))
  return x.sum()

def do(expr):
  eval(expr)
  print(f'{expr}\n\n{cool.pop()}')
  print('\n===\n')
  jax.clear_caches()

x = jnp.ones((3, 3))

do('jax.jit(f)(x)')
do('jax.grad(f)(x)')
do('jax.grad(lambda x: jax.grad(f)(x).sum())(x)')
do('jax.grad(jax.jit(f))(x)')
do('jax.jit(jax.grad(f))(x)')

instead of getting this like we do at HEAD:

jax.jit(f)(x)

Traced<ShapedArray(float32[3,3])>with<DynamicJaxprTrace(level=1/0)>

===

jax.grad(f)(x)

Traced<ConcreteArray([[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]], dtype=float32)>with<JVPTrace(level=2/0)> with
  primal = Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)
  tangent = Traced<ShapedArray(float32[3,3])>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[3,3]), None)
    recipe = LambdaBinding()

===

jax.grad(lambda x: jax.grad(f)(x).sum())(x)

Traced<ConcreteArray([[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]], dtype=float32)>with<JVPTrace(level=4/0)> with
  primal = Traced<ConcreteArray([[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]], dtype=float32)>with<JVPTrace(level=2/0)> with
    primal = Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)
    tangent = Traced<ShapedArray(float32[3,3])>with<JaxprTrace(level=1/0)> with
      pval = (ShapedArray(float32[3,3]), None)
      recipe = LambdaBinding()
  tangent = Traced<ShapedArray(float32[3,3])>with<JaxprTrace(level=3/0)> with
    pval = (ShapedArray(float32[3,3]), None)
    recipe = LambdaBinding()

===

jax.grad(jax.jit(f))(x)

Traced<ShapedArray(float32[3,3])>with<DynamicJaxprTrace(level=3/0)>

===

jax.jit(jax.grad(f))(x)

Traced<ShapedArray(float32[3,3])>with<JVPTrace(level=3/0)> with
  primal = Traced<ShapedArray(float32[3,3])>with<DynamicJaxprTrace(level=1/0)>
  tangent = Traced<ShapedArray(float32[3,3])>with<JaxprTrace(level=2/0)> with
    pval = (ShapedArray(float32[3,3]), None)
    recipe = LambdaBinding()

===

We get this:

jax.jit(f)(x)

JitTracer(float32[3,3])

===

jax.grad(f)(x)

GradTracer(Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32))

===

jax.grad(lambda x: jax.grad(f)(x).sum())(x)

GradTracer(GradTracer(Array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)))

===

jax.grad(jax.jit(f))(x)

JitTracer(float32[3,3])

===

jax.jit(jax.grad(f))(x)

GradTracer(JitTracer(float32[3,3]))

===

I think for this line in the docs we'd get JitTracer(int32[5]) instead of the current Traced<ShapedArray(int32[5])>with<DynamicJaxprTrace(level=1/0)>.

Landing this is outside the scope of this PR review, but I wanted to jot this down for posterity so we don't forget. LMK if you have any thoughts.

[mattjj]

In key-concepts.md there's a line "short for JAX eXPRession" that should be "short for JAX exPRession" to make the capitals work out.

How about we change that sentence from "A jaxpr (short for JAX eXPRession) represents a list of core units of computation called primitives that represent the effect of a computation." to something like "A jaxpr (short for JAX exPRession) is a simple representation of a functional program, comprising a sequence of primitive operations."

[mattjj]

In jit-compilation.md, I think this paragraph might make it sound like global_list.append has more of a chance at predictable behavior than it really does:

Of course, impure functions can still be written and even run, but JAX gives no guarantees about their behaviour under transformations. However, as a rule of thumb, you can expect (but shouldn’t rely on) the side-effects of a JIT-compiled function to run once (during the first call), and never again, due to JAX’s traced execution model.

The trouble is that in the example, we'll leak a tracer, and all heck will break loose!

Maybe we can change it to something like:

Impure functions are dangerous because under JAX transformations they are likely not to behave as intended; they might fail silently, or produce surprising downstream errors like leaked Tracers. Moreover, JAX often can't detect when side effects are present. (If you want debug printing, use jax.debug.print (link to debugging section of tutorial!). To express general side-effects at the cost of performance, see jax.experimental.io_callback (link to docs). To check for tracer leaks at the cost of performance, use with jax.check_tracer_leaks (link to docs).)

I'm not sure about all those forward pointers but they may be useful. I've liked the other forward pointers in here so far!

[mattjj]

In jit-compilation.md, another possible forward-pointer: when we talk about print, we can say it's a trace-time print, and if you want a run-time debug print that does appear in the jaxpr, check out jax.debug.print.

Nevermind, that's exactly what is already said in debugging.md. So a forward pointer covers that.

[mattjj]

This list item indentation in jit-compilation.md is a bit off:

[mattjj]

In jit-compilation.md, I think this section is misleading:

We forgot the block_until-ready on the second example!

I think the whole section is misleading, though, because jit likely always worth it. It should have lower overhead on each execution, and we have to compile in both cases, just the jit-decorated case compiles one larger program rather than multiple smaller ones.

Maybe we can just excise the whole section? The message at the end, to jit the largest chunks of your program that you can, is a good one, so maybe we can just keep that last paragraph?

[mattjj]

In automatic-differentiation.md, we write "automatic differentiation (autodiff)" twice (once in each of the first two sections). The second time could just say "autodiff".

[jakevdp]

In quickstart.md there's a sentence that says "For more advanced autodiff operations, you can use jax..jacrev() for reverse-mode vector-Jacobian products, and jax.jacfwd() for forward-mode Jacobian-vector products"

It previously said something like this, but I found it confusing because the sentence mentions jax.jvp and jax.vjp, then gives an example using jax.jacfwd and jax.jacrev. I fixed this by changing the functions mentioned in the paragraph.

I'm fine with switching back, but that re-introduces the strangeness of the example not using the functions mentioned in the test. What do you suggest?

[jakevdp]

I'm going to close this PR, becuase it's about moving camera-ready content to the right place. The volume of content-related comments is telling me it's not actually camera-ready! It will be cleanest to address those separately in-place, and then move the content once it's ready.

[mattjj]

I'm fine with switching back, but that re-introduces the strangeness of the example not using the functions mentioned in the test. What do you suggest?

Yeah tricky call. Do you think it's okay to use vmap here (with a forward pointer), even though we haven't covered it in the tutorial order yet? If so then I'll try inlining simplified (no pytree handling) definitions of jacfwd and jacrev, so the example looks like vmap + jvp + vmap + vjp = hessian, and we can take a look at that version. Or if it's a priori obvious that we shouldn't mention vmap here, I'm not sure what to do... maybe move it to the vmap section and have a forward pointer here?

[jakevdp]

I don't think we should confuse things by adding vmap. If you want to mention jvp and vjp, then the example should use those. Or, if you want to keep the example as-is, maybe one more sentence explaining why it's relevant would be enough.