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

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

    Returns:
      The output of the final sublayer.
    """
    for i, layer in enumerate(self.sublayers):
      with jax.named_scope(f"{i}"):
        value = layer(value, **side_inputs)
    return value
Runs each of the sublayers in sequence.

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

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