# from jax import tree_util
# from meent.on_jax.optimizer.optimizer import OptimizerJax
#
# tree_util.register_pytree_node(OptimizerJax,
#                                OptimizerJax._tree_flatten,
#                                OptimizerJax._tree_unflatten)
