add error checks for refs, behind a flag

[mattjj]

This PR adds some error checks for incorrect ustage of mutable array references. The checks are behind a flag, default off, because there are probably some downstream users that need to be updated, and we want to defer doing that until all the checks land.

This PR only adds checks for jit, namely:

1. that we don't pass in the same mutable array argument twice;
2. that we don't pass in a mutable array as an argument that we also close over;
3. that we don't return any mutable arrays.

We need to do the same for all higher-order primitives (HOPs): scan, cond, while_loop, and custom_jvp/vjp. I'm going to leave those as follow-ups so we can iterate on getting the jit errors right, since those will largely follow the same pattern.

Here are some example error messages:

import jax
import jax.numpy as jnp
from jax._src import core

jax.config.update('jax_mutable_array_checks', True)

@jax.jit
def f(x):
  return core.mutable_array(x)

f(jnp.arange(3.))

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/dougal70.py", line 11, in <module>
    f(jnp.arange(3.))
ValueError: function f at
/usr/local/google/home/mattjj/packages/jax/dougal70.py:7 traced for jit returned
a mutable array reference of type Ref{float32[3]}, but mutable array references
cannot be returned.

The returned mutable array was created on line 
/usr/local/google/home/mattjj/packages/jax/dougal70.py:9:9 (f).

---

@jax.jit
def f(x):
  return {'hi': [0, core.mutable_array(x), 5]}

f(jnp.arange(3.))

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/dougal70.py", line 11, in <module>
    f(jnp.arange(3.))
ValueError: function f at
/usr/local/google/home/mattjj/packages/jax/dougal70.py:7 traced for jit returned
a mutable array reference of type Ref{float32[3]} at output tree path ['hi'][1],
but mutable array references cannot be returned.

The returned mutable array was created on line 
/usr/local/google/home/mattjj/packages/jax/dougal70.py:9:20 (f).

---

@jax.jit
def f(y_ref):
  return y_ref

x_ref = core.mutable_array(x)
f(x_ref)

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/dougal70.py", line 12, in <module>
    f(x_ref)
ValueError: function f at
/usr/local/google/home/mattjj/packages/jax/dougal70.py:7 traced for jit returned
a mutable array reference of type Ref{float32[3]}, but mutable array references
cannot be returned.

The returned mutable array was passed in as the argument y_ref.

---

@jax.jit
def f(x_ref, y_ref):
  ...

x_ref = core.mutable_array(jnp.arange(3.))
f(x_ref, x_ref)

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/dougal70.py", line 12, in <module>
    f(x_ref, x_ref)
ValueError: only one reference to a mutable array may be passed as an argument
to a function, but when tracing f at
/usr/local/google/home/mattjj/packages/jax/dougal70.py:7 for jit the mutable
array reference of type Ref{float32[3]} appeared at both x_ref and y_ref.

---

@jax.jit
def f(params):
  ...

x_ref = core.mutable_array(jnp.arange(3.))
f({'a': x_ref, 'b': {'c': x_ref}})

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/dougal70.py", line 12, in <module>
    f({'a': x_ref, 'b': {'c': x_ref}})
ValueError: only one reference to a mutable array may be passed as an argument
to a function, but when tracing f at
/usr/local/google/home/mattjj/packages/jax/dougal70.py:7 for jit the mutable
array reference of type Ref{float32[3]} appeared at both params['a'] and
params['b']['c'].