With this change, a value `x` can be replicated `nrep` times as follows:
```python
pmap(lambda _: x)(np.arange(nrep))
```
This will broadcast `x` into a ShardedDeviceArray suitable for passing into another pmap with the same input shape.
If `x` will be passed into a pmap with `devices` or a nested pmap, the replication pmap(s) should follow that structure. For example:
```python
x = pmap(pmap(lambda _: x))(np.ones(2, 4))
pmap(pmap(lambda i: i**2 + x))(np.ones(2, 4))
```
* Support IntEnum values as arguments to JAX functions.
When abstractifying a Python value, search the method-resolution order (MRO) of the type rather than only looking at the value's own type. IntEnum instances are subclasses of int, so this allows us to correctly handle them as integers, much as NumPy itself does.
Testing is done by running "jupyter nbconvert --to notebook" and
then parsing the resulting notebook to look for errors.
One can declare expected errors, and the test will fail if those
are missing.
In the process of doig this, found and fixed a bug in the autodiff_cookbook
notebook.
Also enable JAX_ENABLE_X64=True tests on Travis to avoid future such
issues. (Internal tests catch things like this, but we don't run those
automatically.)
