def METHOD_NAME(self, query, key):
    """Calculates attention scores as a nonlinear sum of query and key.
    Args:
        query: Query tensor of shape `(batch_size, Tq, dim)`.
        key: Key tensor of shape `(batch_size, Tv, dim)`.
    Returns:
        Tensor of shape `(batch_size, Tq, Tv)`.
    """
    # Reshape tensors to enable broadcasting.
    # Reshape into [batch_size, Tq, 1, dim].
    q_reshaped = ops.expand_dims(query, axis=-2)
    # Reshape into [batch_size, 1, Tv, dim].
    k_reshaped = ops.expand_dims(key, axis=-3)
    scale = self.scale if self.use_scale else 1.0
    return ops.sum(scale * ops.tanh(q_reshaped + k_reshaped), axis=-1)