make tracers tree-pretty-print their contents

[mattjj]

Fixes #1476.

import jax.numpy as np
from jax import vmap

def f(x):
  y = x ** 2
  print(y)
  return y

vmap(f)(np.arange(4))

Before this PR, that prints:

Traced<ShapedArray(int32[])>with<BatchTrace(level=0/0)>

After this PR, it prints:

Traced<ShapedArray(int32[])>with<BatchTrace(level=0/0)>
  with val = DeviceArray([0, 1, 4, 9], dtype=int32)
       batch_dim = 0

Here's another example:

import jax.numpy as np
from jax import vmap, jacfwd

def f(x):
  y = x ** 2
  print(y)
  return y

jacfwd(f)(np.arange(4.))

Traced<ConcreteArray([0. 1. 4. 9.])>with<JVPTrace(level=1/0)>
  with primal = DeviceArray([0., 1., 4., 9.], dtype=float32)
       tangent = Traced<ShapedArray(float32[4])>with<BatchTrace(level=0/0)>
                   with val = DeviceArray([[0., 0., 0., 0.],
                                           [0., 2., 0., 0.],
                                           [0., 0., 4., 0.],
                                           [0., 0., 0., 6.]], dtype=float32)
                        batch_dim = 0

The same works in pdb.

One more example!

import jax.numpy as np
from jax import vmap, jacfwd, grad, jit


def f(x):
  y = x ** 2
  print(y)
  return y

grad(f)(3.)
grad(jit(f))(3.)

Traced<ConcreteArray(9.0)>with<JVPTrace(level=1/0)>
  with primal = DeviceArray(9., dtype=float32)
       tangent = Traced<ShapedArray(float32[]):JaxprTrace(level=0/0)>

Traced<ShapedArray(float32[])>with<JVPTrace(level=1/1)>
  with primal = Traced<ShapedArray(float32[]):JaxprTrace(level=-1/1)>
       tangent = Traced<ShapedArray(float32[]):JaxprTrace(level=0/1)>

That shows how grad specializes on concrete primal values, but traces the tangents with abstract values, unless of course there's a jit involved staging out the primal stuff too.

JAX is pretty special in that it performs many transformations on-the-fly while it's evaluating user code, rather than first staging things out of Python into some IR and then doing the transformations on that IR. We can take advantage of that to give users better debugging workflows, using standard Python tools (like prints and pdb).

This PR is just a change to the Tracer repr method.