allow `sow` on big ints (and non jax types), and improve error

Currently this gives a bad error:

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
class Einselm(nn.Module):
@nn.compact
def __call__(self, x):
self.sow('intermediates', 'features', 2 ** 32)
return x

cnn = nn.checkpoint(Einselm)()
variables = cnn.init(jax.random.PRNGKey(0), jnp.ones((1, 28, 28, 1)))
_, sown = cnn.apply(variables, jnp.ones((1, 28, 28, 1)), mutable='intermediates')
```

```
...
File "/usr/local/google/home/mattjj/packages/flax/flax/linen/module.py", line 2061, in scope_fn
return fn(module.clone(parent=scope), *args, **kwargs)
File "/usr/local/google/home/mattjj/packages/flax/flax/linen/transforms.py", line 324, in wrapped_fn
ret = trafo_fn(module_scopes, *args, **kwargs)
File "/usr/local/google/home/mattjj/packages/flax/flax/core/lift.py", line 214, in wrapper
y, out_variable_groups_xs_t = fn(
File "/usr/local/google/home/mattjj/packages/flax/flax/core/lift.py", line 1220, in inner
return rematted(variable_groups, rng_groups, *args, **kwargs)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/ad_checkpoint.py", line 285, in fun_remat
jaxpr, consts, out_tree = _trace_to_jaxpr(fun_, in_tree, tuple(in_avals))
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/ad_checkpoint.py", line 375, in _trace_to_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals, debug)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/interpreters/partial_eval.py", line 2150, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/interpreters/partial_eval.py", line 2173, in trace_to_subjaxpr_dynamic
out_tracers = map(trace.full_raise, ans)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/util.py", line 109, in safe_map
return list(map(f, *args))
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/core.py", line 475, in full_raise
return self.pure(val)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/interpreters/partial_eval.py", line 1798, in new_const
aval = raise_to_shaped(get_aval(c), weak_type=dtypes.is_weakly_typed(c))
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/core.py", line 1334, in get_aval
return concrete_aval(x)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/core.py", line 1323, in concrete_aval
if handler: return handler(x)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/abstract_arrays.py", line 80, in _make_concrete_python_scalar
dtype = dtypes._scalar_type_to_dtype(t, x)
File "/usr/local/google/home/mattjj/packages/jax/jax/_src/dtypes.py", line 202, in _scalar_type_to_dtype
raise OverflowError(f"Python int {value} too large to convert to {dtype}")
OverflowError: Python int 4294967296 too large to convert to int32
```

The error stems from calling `sow` on a big int, and then that value getting plumbed as an output from the `jax.checkpoipnt`, though that's not clear from the traceback!

We'd like to do two things:
1. improve the error message, and
2. give users a way to express what they want without getting an error.

To improve the error message, we could tree-map` jax._src.core.get_aval` or similar on the sown value to check early that the value is a (pytree of) valid jax type and that it doesn't trigger any overflow errors.

To give users a way to express this kind of thing, we could build into flax's `sow` something like this:

```python
from typing import Any

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.struct import dataclass, field

@dataclass
class Static:
val: Any = field(pytree_node=False)

class Einselm(nn.Module):
@nn.compact
def __call__(self, x):
self.sow('intermediates', 'features', Static(2 ** 32))
return x

cnn = nn.checkpoint(Einselm)()
variables = cnn.init(jax.random.PRNGKey(0), jnp.ones((1, 28, 28, 1)))
_, sown = cnn.apply(variables, jnp.ones((1, 28, 28, 1)), mutable='intermediates')

sown = jax.tree_util.tree_map(lambda x: x.val if type(x) is Static else x,
sown, is_leaf=lambda l: type(l) is Static)
print(sown)
```

That is, we're just putting the value in a static field (hopefully it's hashable... we may need to wrap it in a hash-by-id instance if not) and then unboxing it on the output.

cc @levskaya