code: def __init_subclass__(cls):
      super().__init_subclass__()
  
      jtu.register_pytree_with_keys(
          cls,
          partial(_variable_flatten, with_keys=True),  # type: ignore
          partial(_variable_unflatten, cls=cls),  # type: ignore
          flatten_func=partial(_variable_flatten, with_keys=False),  # type: ignore
      )
  
imports: import dataclasses
  import functools
  import typing
  from abc import ABCMeta
  from functools import partial
  from typing import Any
  import jax
  import jax.tree_util
  from flax.experimental.nnx.nnx import reprlib
label_desc: Use image node_id 17 for calling the Variable obj's underlying member method code with example usage: obj.__init_subclass__(cls) without return types