def checked_layer_call(
    func: Callable[[Layer, Any], Any],
) -> Callable[[Layer, Any], Any]:
  """Decorator for ``Layer.__call__`` to add shape-checking and name scopes."""

  @functools.wraps(func)
  def wrapper(self, argument: Any, /) -> Any:
    input_vars = shapecheck.check_structure(
        argument,
        pattern=self.input_structure(),
        error_prefix=(
            "Error while checking the input of a layer of type"
            f" {type(self).__qualname__}:\n"
        ),
    )
    output_structure = self.output_structure()
    with jax.named_scope(type(self).__name__):
      result = func(self, argument)
    shapecheck.check_structure(
        result,
        pattern=output_structure,
        known_vars=input_vars,
        error_prefix=(
            "Error while checking the output of a layer of type"
            f" {type(self).__qualname__}:\n"
        ),
    )
    return result

  wrapper.__penzai_checked_layer_call__ = True
  return wrapper
Decorator for ``Layer.__call__`` to add shape-checking and name scopes.
penzai/deprecated/v1/core/layer.py