Error with Pydantic

[kaiwenw]

Hello! I'm interested in using pydantic's recursive constructor / asdict functionality, but jax.jit-ed functions give the following error:

Argument '_Pydantic_OptimConfig_93971134241088(.. SOMETHING HERE...)' of type <class 'pydantic.dataclasses._Pydantic_OptimConfig_93971134241088'> is not a valid JAX type.