def __call__(self, value: Any, **side_inputs: dict[Any, Any]) -> Any:
    """Runs each of the sublayers in sequence.

    Args:
      value: The input to the layer.
      **side_inputs: The side inputs for all sublayers.

    Returns:
      The output of the final sublayer.
    """
    dimvars = shapecheck.check_structure(
        value,
        self.input_like,
        error_prefix="While checking the input of a CheckedSequential:\n",
    )
    for i, layer in enumerate(self.sublayers):
      with jax.named_scope(f"{i}"):
        value = layer(value, **side_inputs)
    shapecheck.check_structure(
        value,
        self.output_like,
        known_vars=dimvars,
        error_prefix="While checking the output of a CheckedSequential:\n",
    )
    return value
Runs each of the sublayers in sequence.

    Args:
      value: The input to the layer.
      **side_inputs: The side inputs for all sublayers.

    Returns:
      The output of the final sublayer.
    
penzai/nn/grouping.py