def tree_flatten_with_keys(
      self,
  ) -> tuple[Sequence[tuple[Any, Any]], Any]:
    """Flattens this tree node with keys.

    See `jax.tree_util.register_pytree_with_keys_class`.

    This method should not be overridden by subclasses, since
    struct-manipulation code should be able to rely on this implementation (and
    in particular, on ``key_for_field`` producing the JAX keypath keys for each
    field). If you must override this for an advanced use case, consider using
    `pytree_dataclass` without subclassing `Struct`.

    Returns:
      ``(key, child)`` pairs for the node, along with static metadata.
    """
    child_field_names = []
    children = []
    static_fields = {}
    # `self` is guaranteed to be a dataclass by AbstractStructMetaclass
    for field_ in dataclasses.fields(self):  # pytype: disable=wrong-arg-types
      value = getattr(self, field_.name)
      if not is_pytree_node_field(field_):
        static_fields[field_.name] = value
      else:
        child_field_names.append(field_.name)
        children.append((
            self.key_for_field(field_.name),
            value,
        ))
    return children, StructStaticMetadata(child_field_names, static_fields)
Flattens this tree node with keys.

    See `jax.tree_util.register_pytree_with_keys_class`.

    This method should not be overridden by subclasses, since
    struct-manipulation code should be able to rely on this implementation (and
    in particular, on ``key_for_field`` producing the JAX keypath keys for each
    field). If you must override this for an advanced use case, consider using
    `pytree_dataclass` without subclassing `Struct`.

    Returns:
      ``(key, child)`` pairs for the node, along with static metadata.
    
penzai/core/struct.py