def __call__(
      self, x: named_axes.NamedArray, **side_inputs: Any
  ) -> named_axes.NamedArray:
    """Applies the attention mask to the input array.

    Args:
      x: The input array to mask. Usually the matrix of query-key dot products.
      **side_inputs: Side inputs. Must include ``query_positions_input_name``
        and ``kv_positions_input_name``.

    Returns:
      An adjusted matrix of logits, where any value where the mask is False has
      been replaced with the `masked_out_value` argument.
    """
    query_positions = side_inputs[self.query_positions_input_name]
    kv_positions = (
        side_inputs[self.kv_positions_input_name]
        .untag(self.seq_axis)
        .tag(self.kv_seq_axis)
    )
    mask = (query_positions >= kv_positions) & (kv_positions >= 0)
    return named_axes.nmap(jnp.where)(mask, x, self.masked_out_value)
Applies the attention mask to the input array.

    Args:
      x: The input array to mask. Usually the matrix of query-key dot products.
      **side_inputs: Side inputs. Must include ``query_positions_input_name``
        and ``kv_positions_input_name``.

    Returns:
      An adjusted matrix of logits, where any value where the mask is False has
      been replaced with the `masked_out_value` argument.
    
penzai/nn/attention.py