[mutable-arrays] allow internal-only refs in remat

[mattjj]

A few ingredients were needed:

1. in core.py, mark InternalMutableArrayEffect as allowed under remat (just whitelisting it)
2. remove a stale NotImplementedError in the scan discharge rule, solved when we used the aptly-named state_discharge2
3. adapt stateful ops remat partial eval rules, basically to do full remat regardless of remat policy
4. add such a rule for pallas core_map

@dougalm observed a downside to item 3: it means remat with an everything-saveable policy is not the same as AD classic (without the remat decorator). That's because we always remat these internally-stateful operations, regardless of policy. An alternative is to always save such computations (regardless of policy), but then we can't recover full remat. We don't want to risk splitting up stateful operations across the two sides, and it's sufficient to always put them one way or the other (regardless of policy). One way around this would be to statically switch on everything-saveable and/or nothing-saveable policies, represented as something other than an opaque callback.

We still can't handle this code until we redo remat again on top of direct-linearize (remat3):

import jax
import jax.numpy as jnp
jax.config.update('jax_enable_checks', True)

@jax.remat
def f(x):
  x_ref = jax.array_ref(0.)
  x_ref[...] = x
  return x_ref[...]

jax.grad(f)(3.)