Issues checkpointing optimizer state using Optax, nnx.Optimizer, and Orbax

I am running into an error while trying to checkpoint an Optax optimizer state, wrapped as an `nnx.Optimizer`, using the Orbax checkpointing library.

```
ValueError: Unsupported type: for key: ('0', 'count'). Supported types are (, , , ).
```

I am using packages:
- flax 0.10.2
- jax 0.4.36
- jax-cuda12-pjrt 0.4.36
- jax-cuda12-plugin 0.4.36
- jaxlib 0.4.36
- optax 0.2.4
- orbax-checkpoint 0.10.2

Minimal repro:
```
from flax import nnx
import numpy as np
import orbax.checkpoint as ocp
import optax
import os
import pathlib

class MyModel(nnx.Module):
def __init__(self, rngs):
self.linear = nnx.Linear(in_features=4, out_features=1, rngs=rngs)

def __call__(self, x):
return self.linear(x)

rngs = nnx.Rngs(0)
model = MyModel(rngs)

tx = optax.adam(1e-3)
optimizer = nnx.Optimizer(model, tx)

checkpointDir = pathlib.Path('/tmp/my-checkpoints/')
checkpointer = ocp.StandardCheckpointer()

checkpointer.save(checkpointDir, optimizer.opt_state, force=True)
```

I see yall have documentation about using Orbax for model checkpointing, but don't see any official info about optimizer state checkpointing.

I see an older github issue where the Optax folks recommended `cloudpickle`: https://github.com/google-deepmind/optax/discussions/180

I tried adding a custom serialization/deserialization for `nnx.training.optimizer.OptArray` via https://orbax.readthedocs.io/en/latest/guides/checkpoint/custom_handlers.html#custom-serialization-deserialization, but then just ran into the next error:

```
ValueError: TypeHandler lookup failed for: type=
```

Maybe I'd need to also add a type handler for `OptVariable` as well as maybe even `Variable`? That seems quite annoying. Am I missing something?

Thanks for your time!