Преобразуйте большие по модулю логиты в устойчивые вероятности

/// Applies the literal exponential normalization used as a bounded baseline.
///
/// This exposes finite-precision overflow and underflow; it is not attributed
/// to the software implementation of any cited language model.
pub fn direct_output_softmax(logits: &[f64]) -> Vec<f64> {
    let exponentials = logits.iter().map(|value| value.exp()).collect::<Vec<_>>();
    let denominator = exponentials.iter().sum::<f64>();
    exponentials
        .into_iter()
        .map(|value| value / denominator)
        .collect()
}

/// A rejected probability operation, target, output, or converted view operation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbabilityError {
    /// An owned output layout violates the tensor storage invariant.
    Tensor(TensorError),
    /// A tensor-view error was converted into the probability error type.
    View(TensorViewError),
    /// The requested class axis does not exist.
    AxisOutOfBounds { axis: usize, rank: usize },
    /// Softmax, log-softmax, and indexed NLL need at least one class.
    EmptyNormalizationAxis { axis: usize },
    /// The checked output shape is valid, but its value buffer cannot be reserved.
    OutputAllocationFailed { elements: usize },
    /// The first rejected logit in group-major, class-minor order is NaN.
    NaNLogit { group: usize, class: usize },
    /// The first rejected logit in group-major, class-minor order is positive infinity.
    PositiveInfinityLogit { group: usize, class: usize },
    /// The first rejected logit in group-major, class-minor order is negative infinity.
    NegativeInfinityLogit { group: usize, class: usize },
    /// There must be one flat target for every class-axis group.
    TargetCountMismatch { expected: usize, actual: usize },
    /// A mean is undefined when there are no target groups.
    EmptyTargets,
    /// One target does not name a class on the selected axis.
    TargetOutOfBounds {
        group: usize,
        target: usize,
        classes: usize,
    },
}

impl fmt::Display for ProbabilityError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tensor(error) => error.fmt(formatter),
            Self::View(error) => error.fmt(formatter),
            Self::AxisOutOfBounds { axis, rank } => {
                write!(
                    formatter,
                    "probability axis {axis} is out of bounds for rank {rank}"
                )
            }
            Self::EmptyNormalizationAxis { axis } => {
                write!(formatter, "probability axis {axis} has no classes")
            }
            Self::OutputAllocationFailed { elements } => write!(
                formatter,
                "cannot allocate probability output for {elements} f64 values"
            ),
            Self::NaNLogit { group, class } => {
                write!(formatter, "logit at group {group}, class {class} is NaN")
            }
            Self::PositiveInfinityLogit { group, class } => write!(
                formatter,
                "logit at group {group}, class {class} is positive infinity"
            ),
            Self::NegativeInfinityLogit { group, class } => write!(
                formatter,
                "logit at group {group}, class {class} is negative infinity"
            ),
            Self::TargetCountMismatch { expected, actual } => write!(
                formatter,
                "indexed mean NLL needs {expected} targets, but received {actual}"
            ),
            Self::EmptyTargets => formatter.write_str("indexed mean NLL needs at least one target"),
            Self::TargetOutOfBounds {
                group,
                target,
                classes,
            } => write!(
                formatter,
                "target {target} at group {group} is out of bounds for {classes} classes"
            ),
        }
    }
}

impl Error for ProbabilityError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Tensor(error) => Some(error),
            Self::View(error) => Some(error),
            _ => None,
        }
    }
}

impl From<TensorError> for ProbabilityError {
    fn from(error: TensorError) -> Self {
        Self::Tensor(error)
    }
}

impl From<TensorViewError> for ProbabilityError {
    fn from(error: TensorViewError) -> Self {
        Self::View(error)
    }
}

#[derive(Debug)]
struct AxisPlan {
    axis: usize,
    classes: usize,
    group_shape: Vec<usize>,
    group_strides: Vec<usize>,
    groups: usize,
    class_stride: usize,
}

#[derive(Clone, Copy, Debug)]
struct RowStats {
    maximum: f64,
    shifted_exponential_sum: f64,
    log_shifted_exponential_sum: f64,
}

#[derive(Clone, Copy, Debug)]
struct FiniteLogits;

#[derive(Clone, Copy, Debug)]
enum LogitFiniteness {
    Check,
    Validated(FiniteLogits),
}

impl AxisPlan {
    fn new(
        input: &TensorView<'_>,
        axis: usize,
        allow_empty_axis: bool,
    ) -> Result<Self, ProbabilityError> {
        if axis >= input.rank() {
            return Err(ProbabilityError::AxisOutOfBounds {
                axis,
                rank: input.rank(),
            });
        }

        let classes = input.shape()[axis];
        if classes == 0 && !allow_empty_axis {
            return Err(ProbabilityError::EmptyNormalizationAxis { axis });
        }

        let mut group_shape = input.shape().to_vec();
        group_shape.remove(axis);
        let mut group_strides = input.strides().to_vec();
        let class_stride = group_strides.remove(axis);
        let (_, groups) = checked_row_major_layout(&group_shape)?;
        Ok(Self {
            axis,
            classes,
            group_shape,
            group_strides,
            groups,
            class_stride,
        })
    }

    fn group_offsets(&self, input: &TensorView<'_>) -> StridedOffsets {
        input
            .projected_offsets(&self.group_shape, &self.group_strides, self.groups)
            .expect("a checked probability plan retains valid group-base offsets")
    }

    fn output_group_offsets(&self, output_strides: &[usize], output_len: usize) -> StridedOffsets {
        let mut group_strides = output_strides.to_vec();
        group_strides.remove(self.axis);
        StridedOffsets::checked(
            &self.group_shape,
            &group_strides,
            0,
            self.groups,
            output_len,
        )
        .expect("a checked probability output retains valid group-base offsets")
    }

    fn target_offset(&self, group_base: usize, target: usize) -> usize {
        let class_offset = target
            .checked_mul(self.class_stride)
            .expect("a checked probability plan cannot overflow a class offset");
        group_base
            .checked_add(class_offset)
            .expect("a checked probability plan cannot overflow a target offset")
    }

    fn for_each_group(
        &self,
        input: &TensorView<'_>,
        finiteness: LogitFiniteness,
        mut visit: impl FnMut(usize, usize, RowStats) -> Result<(), ProbabilityError>,
    ) -> Result<(), ProbabilityError> {
        for (group, group_base) in self.group_offsets(input).enumerate() {
            let stats = row_stats(input, self, finiteness, group, group_base)?;
            visit(group, group_base, stats)?;
        }
        Ok(())
    }
}

fn checked_finite_logit(value: f64, group: usize, class: usize) -> Result<f64, ProbabilityError> {
    if value.is_nan() {
        Err(ProbabilityError::NaNLogit { group, class })
    } else if value == f64::INFINITY {
        Err(ProbabilityError::PositiveInfinityLogit { group, class })
    } else if value == f64::NEG_INFINITY {
        Err(ProbabilityError::NegativeInfinityLogit { group, class })
    } else {
        Ok(value)
    }
}

fn validate_finite_logits(
    input: &TensorView<'_>,
    plan: &AxisPlan,
) -> Result<FiniteLogits, ProbabilityError> {
    for (group, group_base) in plan.group_offsets(input).enumerate() {
        let mut input_offset = group_base;
        for class in 0..plan.classes {
            checked_finite_logit(input.value_at_storage_offset(input_offset), group, class)?;
            if class + 1 < plan.classes {
                input_offset = input_offset
                    .checked_add(plan.class_stride)
                    .expect("a checked probability plan cannot overflow along the class axis");
            }
        }
    }
    Ok(FiniteLogits)
}

fn row_stats(
    input: &TensorView<'_>,
    plan: &AxisPlan,
    finiteness: LogitFiniteness,
    group: usize,
    group_base: usize,
) -> Result<RowStats, ProbabilityError> {
    debug_assert!(plan.classes > 0);
    let mut maximum = f64::NEG_INFINITY;
    let mut input_offset = group_base;
    for class in 0..plan.classes {
        let value = input.value_at_storage_offset(input_offset);
        let value = match finiteness {
            LogitFiniteness::Check => checked_finite_logit(value, group, class)?,
            LogitFiniteness::Validated(_) => value,
        };
        maximum = maximum.max(value);
        if class + 1 < plan.classes {
            input_offset = input_offset
                .checked_add(plan.class_stride)
                .expect("a checked probability plan cannot overflow along the class axis");
        }
    }

    let mut exponential_tail = 0.0;
    let mut skipped_one_maximum = false;
    input_offset = group_base;
    for class in 0..plan.classes {
        let value = input.value_at_storage_offset(input_offset);
        let shifted = value - maximum;
        if shifted == 0.0 && !skipped_one_maximum {
            skipped_one_maximum = true;
        } else {
            exponential_tail += shifted.exp();
        }
        if class + 1 < plan.classes {
            input_offset = input_offset
                .checked_add(plan.class_stride)
                .expect("a checked probability plan cannot overflow along the class axis");
        }
    }
    debug_assert!(skipped_one_maximum);

    Ok(RowStats {
        maximum,
        shifted_exponential_sum: 1.0 + exponential_tail,
        log_shifted_exponential_sum: exponential_tail.ln_1p(),
    })
}

fn output_buffer(elements: usize) -> Result<Vec<f64>, ProbabilityError> {
    let mut values = Vec::new();
    values
        .try_reserve_exact(elements)
        .map_err(|_| ProbabilityError::OutputAllocationFailed { elements })?;
    values.resize(elements, 0.0);
    Ok(values)
}

fn positive_zero(value: f64) -> f64 {
    if value == 0.0 { 0.0 } else { value }
}

/// Requested normalized values emitted by one checked forward traversal.
#[derive(Debug)]
struct NormalizedForward {
    probabilities: Option<Tensor>,
    log_probabilities: Option<Tensor>,
}

/// Log-softmax output and optional probabilities from the same forward traversal.
#[derive(Debug)]
pub(crate) struct LogSoftmaxForward {
    pub(crate) value: Tensor,
    pub(crate) probabilities: Option<Tensor>,
}

/// Indexed mean NLL and optional probabilities emitted by its forward traversal.
#[derive(Debug)]
pub(crate) struct IndexedMeanNllForward {
    pub(crate) loss: f64,
    pub(crate) probabilities: Option<Tensor>,
}

struct NormalizedGroupOutput<'a> {
    output_group_base: usize,
    output_class_stride: usize,
    probabilities: Option<&'a mut [f64]>,
    log_probabilities: Option<&'a mut [f64]>,
}

impl NormalizedGroupOutput<'_> {
    fn emit(
        mut self,
        input: &TensorView<'_>,
        plan: &AxisPlan,
        input_group_base: usize,
        stats: RowStats,
    ) {
        debug_assert!(self.probabilities.is_some() || self.log_probabilities.is_some());
        let mut input_offset = input_group_base;
        let mut output_offset = self.output_group_base;
        for class in 0..plan.classes {
            let shifted = input.value_at_storage_offset(input_offset) - stats.maximum;
            if let Some(values) = self.probabilities.as_mut() {
                values[output_offset] =
                    positive_zero(shifted.exp() / stats.shifted_exponential_sum);
            }
            if let Some(values) = self.log_probabilities.as_mut() {
                values[output_offset] = positive_zero(shifted - stats.log_shifted_exponential_sum);
            }

            if class + 1 < plan.classes {
                input_offset = input_offset
                    .checked_add(plan.class_stride)
                    .expect("a checked probability plan cannot overflow along the class axis");
                output_offset = output_offset
                    .checked_add(self.output_class_stride)
                    .expect("a checked probability output cannot overflow along the class axis");
            }
        }
    }
}

/// Reduces one axis with max-shifted log-sum-exp.
///
/// An empty selected axis returns the log-additive identity, negative infinity,
/// once per remaining-axis group. Other non-finite logits are rejected in
/// group-major, class-minor order.
pub fn log_sum_exp(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
) -> Result<Tensor, ProbabilityError> {
    let plan = AxisPlan::new(input, axis, true)?;
    let output_shape = if keep_dim {
        let mut shape = input.shape().to_vec();
        shape[axis] = 1;
        shape
    } else {
        plan.group_shape.clone()
    };
    let (_, output_len) = checked_row_major_layout(&output_shape)?;
    debug_assert_eq!(output_len, plan.groups);
    let mut values = output_buffer(output_len)?;

    if plan.classes == 0 {
        values.fill(f64::NEG_INFINITY);
    } else {
        plan.for_each_group(
            input,
            LogitFiniteness::Check,
            |group, _group_base, stats| {
                values[group] = stats.maximum + stats.log_shifted_exponential_sum;
                Ok(())
            },
        )?;
    }

    Tensor::from_vec(output_shape, values).map_err(Into::into)
}

/// Converts finite logits to normalized probabilities along one explicit axis.
pub fn softmax(input: &TensorView<'_>, axis: usize) -> Result<Tensor, ProbabilityError> {
    let forward = normalized_forward(input, axis, true, false)?;
    Ok(forward
        .probabilities
        .expect("softmax requests a probability output"))
}

/// Converts finite logits to normalized log-probabilities along one explicit axis.
pub fn log_softmax(input: &TensorView<'_>, axis: usize) -> Result<Tensor, ProbabilityError> {
    let forward = log_softmax_forward(input, axis, false)?;
    debug_assert!(forward.probabilities.is_none());
    Ok(forward.value)
}

pub(crate) fn log_softmax_forward(
    input: &TensorView<'_>,
    axis: usize,
    emit_probabilities: bool,
) -> Result<LogSoftmaxForward, ProbabilityError> {
    let forward = normalized_forward(input, axis, emit_probabilities, true)?;
    Ok(LogSoftmaxForward {
        value: forward
            .log_probabilities
            .expect("log-softmax forward requests a log-probability output"),
        probabilities: forward.probabilities,
    })
}

fn normalized_forward(
    input: &TensorView<'_>,
    axis: usize,
    emit_probabilities: bool,
    emit_log_probabilities: bool,
) -> Result<NormalizedForward, ProbabilityError> {
    debug_assert!(emit_probabilities || emit_log_probabilities);
    let plan = AxisPlan::new(input, axis, false)?;
    let (output_strides, output_len) = checked_row_major_layout(input.shape())?;
    let mut log_probability_values = emit_log_probabilities
        .then(|| output_buffer(output_len))
        .transpose()?;
    let finiteness = if emit_probabilities && emit_log_probabilities {
        LogitFiniteness::Validated(validate_finite_logits(input, &plan)?)
    } else {
        LogitFiniteness::Check
    };
    let mut probability_values = emit_probabilities
        .then(|| output_buffer(output_len))
        .transpose()?;
    let output_class_stride = output_strides[axis];

    let mut output_group_offsets = plan.output_group_offsets(&output_strides, output_len);
    plan.for_each_group(input, finiteness, |_group, input_group_base, stats| {
        let output_group_base = output_group_offsets
            .next()
            .expect("a checked probability output has one base per input group");
        NormalizedGroupOutput {
            output_group_base,
            output_class_stride,
            probabilities: probability_values.as_deref_mut(),
            log_probabilities: log_probability_values.as_deref_mut(),
        }
        .emit(input, &plan, input_group_base, stats);
        Ok(())
    })?;
    debug_assert!(output_group_offsets.next().is_none());

    let log_probabilities = log_probability_values
        .map(|values| Tensor::from_vec(input.shape().to_vec(), values))
        .transpose()?;
    let probabilities = probability_values
        .map(|values| Tensor::from_vec(input.shape().to_vec(), values))
        .transpose()?;
    Ok(NormalizedForward {
        probabilities,
        log_probabilities,
    })
}

/// Scores one class index per remaining-axis group with fused stable mean NLL.
///
/// Targets follow the row-major group shape obtained by removing `axis` from
/// the logits. Bounds are checked for every target before a logit is read.
pub fn indexed_mean_nll(
    logits: &TensorView<'_>,
    axis: usize,
    targets: &[usize],
) -> Result<f64, ProbabilityError> {
    let forward = indexed_mean_nll_forward(logits, axis, targets, false)?;
    debug_assert!(forward.probabilities.is_none());
    Ok(forward.loss)
}

pub(crate) fn indexed_mean_nll_forward(
    logits: &TensorView<'_>,
    axis: usize,
    targets: &[usize],
    emit_probabilities: bool,
) -> Result<IndexedMeanNllForward, ProbabilityError> {
    let plan = AxisPlan::new(logits, axis, false)?;
    if targets.len() != plan.groups {
        return Err(ProbabilityError::TargetCountMismatch {
            expected: plan.groups,
            actual: targets.len(),
        });
    }
    if targets.is_empty() {
        return Err(ProbabilityError::EmptyTargets);
    }
    for (group, &target) in targets.iter().enumerate() {
        if target >= plan.classes {
            return Err(ProbabilityError::TargetOutOfBounds {
                group,
                target,
                classes: plan.classes,
            });
        }
    }

    let finiteness = if emit_probabilities {
        LogitFiniteness::Validated(validate_finite_logits(logits, &plan)?)
    } else {
        LogitFiniteness::Check
    };

    let output_layout = emit_probabilities
        .then(|| checked_row_major_layout(logits.shape()))
        .transpose()?;
    let mut probability_values = output_layout
        .as_ref()
        .map(|(_, output_len)| output_buffer(*output_len))
        .transpose()?;
    let mut output_group_offsets = output_layout
        .as_ref()
        .map(|(output_strides, output_len)| plan.output_group_offsets(output_strides, *output_len));
    let output_class_stride = output_layout
        .as_ref()
        .map(|(output_strides, _)| output_strides[axis]);

    let mut total = 0.0;
    let mut scaled_mean = 0.0;
    let mut needs_scaled_fallback = false;
    let target_count = targets.len() as f64;
    plan.for_each_group(logits, finiteness, |group, group_base, stats| {
        let target = targets[group];
        let target_logit = logits.value_at_storage_offset(plan.target_offset(group_base, target));
        let gap = stats.maximum - target_logit;
        let scaled_gap = if gap.is_finite() {
            gap / target_count
        } else {
            stats.maximum / target_count - target_logit / target_count
        };
        scaled_mean += scaled_gap + stats.log_shifted_exponential_sum / target_count;

        let loss = gap + stats.log_shifted_exponential_sum;
        if loss.is_finite() && !needs_scaled_fallback {
            total += loss;
            if !total.is_finite() {
                needs_scaled_fallback = true;
            }
        } else {
            needs_scaled_fallback = true;
        }

        if let Some(values) = probability_values.as_deref_mut() {
            let output_group_base = output_group_offsets
                .as_mut()
                .and_then(Iterator::next)
                .expect("a checked probability output has one base per input group");
            NormalizedGroupOutput {
                output_group_base,
                output_class_stride: output_class_stride
                    .expect("a requested probability output has a class stride"),
                probabilities: Some(values),
                log_probabilities: None,
            }
            .emit(logits, &plan, group_base, stats);
        }
        Ok(())
    })?;
    debug_assert!(
        output_group_offsets
            .as_mut()
            .is_none_or(|offsets| offsets.next().is_none())
    );

    let loss = positive_zero(if needs_scaled_fallback {
        scaled_mean
    } else {
        total / target_count
    });
    let probabilities = probability_values
        .map(|values| Tensor::from_vec(logits.shape().to_vec(), values))
        .transpose()?;
    Ok(IndexedMeanNllForward {
        loss,
        probabilities,
    })
}

/// Normalizes the same relative logits after three different constant shifts.
pub fn tiny_stable_softmax_example() -> Result<TinyStableSoftmaxExample, ProbabilityError> {
    let logits = Tensor::from_vec(LOGIT_SHAPE.to_vec(), LOGIT_VALUES.to_vec())?;
    let probabilities = softmax(&logits.view(), CLASS_AXIS)?;
    let log_probabilities = log_softmax(&logits.view(), CLASS_AXIS)?;
    let log_normalizers = log_sum_exp(&logits.view(), CLASS_AXIS, false)?;
    let mean_nll = indexed_mean_nll(&logits.view(), CLASS_AXIS, &TARGETS)?;

    Ok(TinyStableSoftmaxExample {
        logits,
        probabilities,
        log_probabilities,
        log_normalizers,
        mean_nll,
    })
}

    let example = tiny_stable_softmax_example()?;
    let row_sums = example
        .probabilities
        .as_slice()
        .chunks_exact(2)
        .map(|row| row.iter().sum())
        .collect::<Vec<f64>>();
    let target_losses = TARGETS
        .iter()
        .enumerate()
        .map(|(row, &target)| -example.log_probabilities.as_slice()[row * 2 + target])
        .collect::<Vec<_>>();
    let ordinary_direct = direct_output_softmax(&[0.0, 1.0]);
    let overflow_direct = direct_output_softmax(&[1000.0, 1001.0]);
    let underflow_direct = direct_output_softmax(&[-1001.0, -1000.0]);

    let axis_error = softmax(&example.logits.view(), 2).unwrap_err();
    let empty_logits = Tensor::from_vec(vec![2, 0], vec![])?;
    let empty_error = softmax(&empty_logits.view(), 1).unwrap_err();
    let nonfinite_logits = Tensor::from_vec(vec![1, 2], vec![0.0, f64::INFINITY])?;
    let nonfinite_error = softmax(&nonfinite_logits.view(), 1).unwrap_err();
    let target_error = indexed_mean_nll(&example.logits.view(), 1, &[1, 2, 1]).unwrap_err();