Method select

penzai/core/struct.py:551–568

Wraps this struct in a selection, enabling functional-style mutations. This is a convenience wrapper around selectors.select to enable easier selection, using syntax like:: struct.select().at(lambda b: b.foo[3].bar).apply(baz) struct.select().at_instances_of(Sequential).at_chil

(self)

  @typing.final
  def select(self) -> selectors.Selection[Struct]:
    """Wraps this struct in a selection, enabling functional-style mutations.

    This is a convenience wrapper around selectors.select to enable easier
    selection, using syntax like::

      struct.select().at(lambda b: b.foo[3].bar).apply(baz)
      struct.select().at_instances_of(Sequential).at_children().apply(qux)

    See documentation for `selectors.Selection` for supported attributes.

    Returns:
      A singleton selection containing this struct.
    """
    # Dynamic import to avoid circular import issues.
    from penzai.core import selectors  # pylint: disable=import-outside-toplevel

    return selectors.select(self)

  def key_for_field(self, field_name: str) -> Hashable:
    """Generates a JAX PyTree key for a given field name.