def key_for_field(self, field_name: str) -> Hashable:
    """Returns a custom GetAttrKey with layer stack metadata."""
    return LayerStackGetAttrKey(
        name=field_name,
        stack_axis=self.stack_axis,
        stack_axis_size=self.stack_axis_size,
    )
Returns a custom GetAttrKey with layer stack metadata.
penzai/nn/layer_stack.py