onnx-export-rs 0.1.1

Export canonical Rust machine-learning models to ONNX
Documentation
use ndarray::{Array1, Array2};

use crate::{Error, Result};

/// A classifier selecting the largest affine class score.
///
/// This captures multinomial and Bernoulli Naive Bayes after their fitted
/// log-probabilities are algebraically reduced to `X * coefficients + bias`.
#[derive(Clone, Debug, PartialEq)]
pub struct LinearScoreClassifier {
    /// Score coefficients shaped `[features, classes]`.
    pub coefficients: Array2<f64>,
    /// One score offset per class.
    pub bias: Array1<f64>,
    /// Integer label corresponding to each score column.
    pub class_labels: Vec<i64>,
    /// Optional strict threshold used to binarize every input feature.
    pub binarize: Option<f64>,
}

impl LinearScoreClassifier {
    /// Creates a validated score classifier.
    ///
    /// # Errors
    ///
    /// Returns an error for inconsistent, empty, duplicate-label, or
    /// non-finite parameters.
    pub fn new(
        coefficients: Array2<f64>,
        bias: Array1<f64>,
        class_labels: Vec<i64>,
        binarize: Option<f64>,
    ) -> Result<Self> {
        let classes = coefficients.ncols();
        let mut unique = class_labels.clone();
        unique.sort_unstable();
        unique.dedup();
        if coefficients.nrows() == 0
            || classes == 0
            || bias.len() != classes
            || class_labels.len() != classes
            || unique.len() != classes
            || coefficients
                .iter()
                .chain(bias.iter())
                .any(|value| !value.is_finite())
            || binarize.is_some_and(|value| !value.is_finite())
        {
            return Err(Error::InvalidModel(
                "invalid linear score classifier".into(),
            ));
        }
        Ok(Self {
            coefficients,
            bias,
            class_labels,
            binarize,
        })
    }
}