onnx-export-rs 0.1.1

use ndarray::{Array1, Array2};

/// Logistic-regression weights for binary or multiclass inference.
///
/// Binary models use exactly one coefficient row and one intercept while
/// setting `n_classes` to two. Their exported output is the positive-class
/// probability with shape `[batch, 1]`. Multiclass models use one row and
/// intercept per class.
#[derive(Clone, Debug, PartialEq)]
pub struct LogisticModelWeights {
    /// Shape `[score_count, feature_count]`.
    pub coefficients: Array2<f64>,
    /// One bias per score.
    pub intercept: Array1<f64>,
    /// Two for binary classification, otherwise the score count.
    pub n_classes: usize,
}

impl LogisticModelWeights {
    /// Constructs logistic model weights, checking the documented convention.
    pub fn new(
        coefficients: Array2<f64>,
        intercept: Array1<f64>,
        n_classes: usize,
    ) -> crate::Result<Self> {
        let scores = coefficients.nrows();
        let valid = n_classes >= 2
            && intercept.len() == scores
            && ((n_classes == 2 && scores == 1) || (n_classes > 2 && scores == n_classes));
        if !valid {
            return Err(crate::Error::InvalidModel(format!(
                "logistic shape mismatch: {scores} coefficient rows, {} intercepts, {n_classes} classes",
                intercept.len()
            )));
        }
        Ok(Self {
            coefficients,
            intercept,
            n_classes,
        })
    }

    /// Returns the required number of input features.
    #[must_use]
    pub fn n_features(&self) -> usize {
        self.coefficients.ncols()
    }
}