code: def __init_subclass__(cls, mutable: bool = False):
      super().__init_subclass__()
      # init class variables
      cls._pytree__is_mutable = False
      cls._pytree__class_is_mutable = mutable
  
      # TODO: clean up this in the future once minimal supported version is 0.4.7
      if hasattr(jax.tree_util, "register_pytree_with_keys"):
          if (
              "flatten_func"
              in inspect.signature(
                  jax.tree_util.register_pytree_with_keys
              ).parameters
          ):
              jax.tree_util.register_pytree_with_keys(
                  cls,
                  partial(
                      cls._pytree__flatten,
                      with_key_paths=True,
                  ),
                  cls._pytree__unflatten,
                  flatten_func=partial(
                      cls._pytree__flatten,
                      with_key_paths=False,
                  ),
              )
          else:
              jax.tree_util.register_pytree_with_keys(
                  cls,
                  partial(
                      cls._pytree__flatten,
                      with_key_paths=True,
                  ),
                  cls._pytree__unflatten,
              )
      else:
          jax.tree_util.register_pytree_node(
              cls,
              partial(
                  cls._pytree__flatten,
                  with_key_paths=False,
              ),
              cls._pytree__unflatten,
          )
  
      # flax serialization support
      if importlib.util.find_spec("flax") is not None:
          from flax import serialization
  
          serialization.register_serialization_state(
              cls, cls._to_flax_state_dict, cls._from_flax_state_dict
          )
  
imports: import contextlib
  import dataclasses
  import importlib.util
  import inspect
  import typing
  from abc import ABCMeta
  from copy import copy
  from functools import partial
  from types import MappingProxyType
  import jax
  import numpy
  from flax.experimental.nnx.nnx import module
  from flax.experimental.nnx.nnx import reprlib, variables
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 2 for calling the Pytree obj's underlying member method code with example usage: obj.__init_subclass__(cls, mutable) without return types