Pickling a trained model (NNX)

I train small models and while prototyping and testing I wish to store trained models in a simple way (also having the model configuration and training-stats inside the model object as dicts/arrays). Idially i want to deal with a single object and being able to simply:

```python
with open('my-model.pk', 'wb') as file:
pickle.dump(model, file)
```
and later:
```python
model = pickle.load('my-model.pk')
```

Similar to `torch.save`.

When i naively try to do the above I get:

```
----> 2 pickle.dump(model, file)
AttributeError: Can't pickle local object 'variance_scaling..init'
```

Splitting into `graphdef` and `state` results in the following error:
```
cannot pickle 'PyTreeDef' object
```

While I am aware there is `orbax` and it might save the state, I really wish it would be possible to avoid that dependency and keep things simple. Is there a trick or a workaround I can use to achieve the desired functionality?