code: def __init__(
      self,
      graphdef: GraphDef[M],
      *,
      params: State,
      tx: optax.GradientTransformation,
      step: int = 0,
      **kwargs,
  ):
      self.graphdef = graphdef
      self.params: State = pytreelib.TreeNode(params)
      self.tx = tx
      self.opt_state = pytreelib.TreeNode(tx.init(self.params))
      self.step = pytreelib.TreeNode(jnp.asarray(step))
      for name, value in kwargs.items():
          if isinstance(value, (jax.Array, np.ndarray, State)):
              value = pytreelib.TreeNode(value)
          setattr(self, name, value)
  
imports: from __future__ import annotations
  import inspect
  import typing
  import jax
  import jax.numpy
  import numpy
  import optax
  from flax.experimental.nnx.nnx import pytreelib
  from flax.experimental.nnx.nnx.module import GraphDef, Module
  from flax.experimental.nnx.nnx.proxy_caller import ApplyCaller
  from flax.experimental.nnx.nnx.rnglib import Rngs
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 1 to create a new TrainState object from inherited base classes: pytreelib, tp.Generic with example: obj = TrainState(graphdef)