def broadcast_to(
      self,
      positional_shape: Sequence[int] = (),
      named_shape: Mapping[AxisName, int] | None = None,
  ) -> NamedArrayBase:
    """Broadcasts a named array to a possibly-larger shape.

    Args:
      positional_shape: Desired positional shape for the array. Will be
        broadcast using numpy broadcasting rules.
      named_shape: Desired named shape for the array. Will be broadcast using
        `nmap`-style vectorizing rules (e.g. new named axes will be introduced
        if missing, but length-1 axes will not be broadcast).

    Returns:
      A named array that has the given positional and named shapes. Note that
      if this array has axis names that are not in ``named_shape``, these will
      be preserved in the answer as well.
    """
    if named_shape is None:
      named_shape = {}
    named_shape = dict(named_shape)
    if (
        self.positional_shape == tuple(positional_shape)
        and dict(self.named_shape) == named_shape
    ):
      return self

    # Trick: create a size-zero array with the right shape so that we can
    # broadcast using nmap's vectorization rules.
    prototype_data = jnp.zeros(
        tuple(named_shape.values()) + tuple(positional_shape) + (0,)
    )
    assert prototype_data.size == 0
    prototype = NamedArray.wrap(prototype_data).tag_prefix(*named_shape.keys())
    return nmap(lambda a, b: jnp.broadcast_to(a, b.shape[:-1]))(self, prototype)
Broadcasts a named array to a possibly-larger shape.

    Args:
      positional_shape: Desired positional shape for the array. Will be
        broadcast using numpy broadcasting rules.
      named_shape: Desired named shape for the array. Will be broadcast using
        `nmap`-style vectorizing rules (e.g. new named axes will be introduced
        if missing, but length-1 axes will not be broadcast).

    Returns:
      A named array that has the given positional and named shapes. Note that
      if this array has axis names that are not in ``named_shape``, these will
      be preserved in the answer as well.
    
penzai/core/named_axes.py