Deepcopy of pjit functions failing (or not supported)

[jlperla]

Description

While trying to get jsonargparse to work with defaults on activation functions, I ran into omni-us/jsonargparse#619 (comment)

The preferred MWE from the developer there is repeated below

>>> from copy import deepcopy
>>> from dataclasses import dataclass, asdict
>>> from flax import nnx
>>> from typing import Callable
>>> @dataclass
... class Settings:
...     width: int = 64
...     activation: Callable = nnx.relu
...     final_activation: Callable = nnx.softplus
>>> deepcopy(nnx.softplus) is nnx.softplus
False
>>> nnx.softplus
<PjitFunction of <function softplus at 0x7d73a9f0fb50>>
>>> deepcopy(nnx.softplus)
<PjitFunction>

It seems like the issue is that a deepcopy and perhaps a asdict of the jit'd functions is failing or not supported.

As said in that post, the workaround shows there is no intrinsic reason deepcopy can't be supported. If we put

from types import MethodType
nnx.softplus.__deepcopy__ = MethodType(lambda self, memo: self, nnx.softplus)

Then things work great.

System info (python version, jaxlib version, accelerator, etc.)

This applies to all versions.