onnx-export-rs 0.1.1

//! Integration tests for graph construction and model exporters.

use ndarray::{array, ArrayD};
use onnx_export_rs::canonical::{
    flatten_tree, AggregationMode, ForestStructure, GeneralizedLinearModel, GradientBoostedEnsemble,
    LinearModelWeights, LinkFunction, LogisticModelWeights, PostTransform, RecursiveNode,
    SvmClassifier, SvmKernel, SvmRegressor, TreeTask,
};
use onnx_export_rs::exporters::{
    export_generalized_linear, export_gradient_boosting, export_linear, export_logistic,
    export_svm_classifier, export_svm_regressor, export_tree_ensemble,
};
use onnx_export_rs::graph_builder::{assemble_model, make_node, make_value_info, Dimension};
use onnx_export_rs::proto::tensor_shape_proto::dimension::Value;
use onnx_export_rs::proto::{GraphProto, ModelProto};
use onnx_export_rs::{to_bytes, IR_VERSION, OPSET_VERSION};
use prost::Message;

#[test]
fn identity_model_encodes_and_decodes() {
    let graph = GraphProto {
        node: vec![make_node("Identity", ["X"], ["Y"], vec![])],
        name: "identity".into(),
        initializer: vec![],
        doc_string: String::new(),
        input: vec![make_value_info(
            "X",
            &[Dimension::Symbolic("batch".into()), Dimension::Fixed(2)],
        )],
        output: vec![make_value_info(
            "Y",
            &[Dimension::Symbolic("batch".into()), Dimension::Fixed(2)],
        )],
        value_info: vec![],
    };
    let bytes = to_bytes(&assemble_model(graph, OPSET_VERSION, IR_VERSION)).unwrap();
    let decoded = ModelProto::decode(bytes.as_slice()).unwrap();
    assert_eq!(decoded.graph.unwrap().node[0].op_type, "Identity");
}

#[test]
fn exports_linear_gemm_with_expected_shapes() {
    let model = export_linear(&LinearModelWeights::new(array![2.0, -1.0, 0.5], 3.0));
    let graph = model.graph.unwrap();
    assert_eq!(graph.node[0].op_type, "Gemm");
    assert_eq!(graph.initializer[0].dims, [3, 1]);
    assert_eq!(graph.initializer[0].float_data, [2.0, -1.0, 0.5]);
}

#[test]
fn exports_generalized_linear_link_activations() {
    let identity =
        export_generalized_linear(&GeneralizedLinearModel::new(array![1.0], 0.0, LinkFunction::Identity));
    let identity_nodes = identity.graph.unwrap().node;
    assert_eq!(identity_nodes.len(), 1);
    assert_eq!(identity_nodes[0].op_type, "Gemm");
    assert_eq!(identity_nodes[0].output, ["Y"]);

    let log = export_generalized_linear(&GeneralizedLinearModel::new(
        array![0.3, -0.1],
        0.5,
        LinkFunction::Log,
    ));
    let log_nodes = log.graph.unwrap().node;
    assert_eq!(log_nodes[0].op_type, "Gemm");
    assert_eq!(log_nodes[0].output, ["scores"]);
    assert_eq!(log_nodes[1].op_type, "Exp");
    assert_eq!(log_nodes[1].output, ["Y"]);

    let logit = export_generalized_linear(&GeneralizedLinearModel::new(
        array![1.0],
        0.0,
        LinkFunction::Logit,
    ));
    assert_eq!(logit.graph.unwrap().node[1].op_type, "Sigmoid");
}

#[cfg(feature = "validate")]
#[test]
fn tract_runs_exported_log_link_glm() {
    use ndarray::{array, Array1, Array2};

    let model = GeneralizedLinearModel::new(array![0.4, -0.2], 0.1, LinkFunction::Log);
    let bytes = to_bytes(&export_generalized_linear(&model)).unwrap();
    let inputs: Array2<f64> = array![[1.0, 2.0], [0.5, -1.0], [-2.0, 3.0]];
    let coefficients = model.coefficients.clone();
    let intercept = model.intercept;
    let report = onnx_export_rs::validate::validate_export(&bytes, &inputs, |batch| {
        Array1::from_iter(
            batch
                .rows()
                .into_iter()
                .map(|row| (row.dot(&coefficients) + intercept).exp()),
        )
    }, 1e-5)
    .unwrap();
    assert!(report.passed, "log-link GLM mismatch: {report:?}");
}

#[test]
fn exports_binary_and_multiclass_activation() {
    let binary = LogisticModelWeights::new(array![[1.0, 2.0]], array![0.0], 2).unwrap();
    assert_eq!(
        export_logistic(&binary).graph.unwrap().node[1].op_type,
        "Sigmoid"
    );
    let multi = LogisticModelWeights::new(
        array![[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]],
        array![0.0, 0.0, 0.0],
        3,
    )
    .unwrap();
    let node = &export_logistic(&multi).graph.unwrap().node[1];
    assert_eq!(node.op_type, "Softmax");
    assert_eq!(node.attribute[0].i, 1);
}

#[test]
fn flattens_and_exports_tree() {
    let tree = flatten_tree(&RecursiveNode::Branch {
        feature: 1,
        threshold: 2.5,
        left: Box::new(RecursiveNode::Leaf(vec![10.0])),
        right: Box::new(RecursiveNode::Leaf(vec![-3.0])),
    })
    .unwrap();
    assert_eq!(tree.nodes.len(), 3);
    assert_eq!(
        (tree.nodes[0].true_child_id, tree.nodes[0].false_child_id),
        (1, 2)
    );
    let forest = ForestStructure {
        trees: vec![tree],
        aggregation: AggregationMode::Average,
        n_targets: 1,
    };
    let model = export_tree_ensemble(&forest, TreeTask::Regression).unwrap();
    let graph = model.graph.unwrap();
    assert_eq!(graph.node[0].domain, "ai.onnx.ml");
    assert_eq!(graph.node[0].op_type, "TreeEnsembleRegressor");
}

#[test]
fn tensor_data_is_row_major() {
    let data: ArrayD<f32> = array![[1.0, 2.0], [3.0, 4.0]].into_dyn();
    let tensor = onnx_export_rs::graph_builder::make_tensor("x", &data);
    assert_eq!(tensor.float_data, [1.0, 2.0, 3.0, 4.0]);
}

#[test]
fn handles_degenerate_and_unbalanced_trees() {
    let root_only = flatten_tree(&RecursiveNode::Leaf(vec![4.25])).unwrap();
    assert_eq!(root_only.nodes.len(), 1);
    assert!(root_only.nodes[0].is_leaf());

    let unbalanced = flatten_tree(&RecursiveNode::Branch {
        feature: 0,
        threshold: 0.0,
        left: Box::new(RecursiveNode::Leaf(vec![1.0])),
        right: Box::new(RecursiveNode::Branch {
            feature: 1,
            threshold: 1.0,
            left: Box::new(RecursiveNode::Leaf(vec![2.0])),
            right: Box::new(RecursiveNode::Branch {
                feature: 2,
                threshold: 2.0,
                left: Box::new(RecursiveNode::Leaf(vec![3.0])),
                right: Box::new(RecursiveNode::Leaf(vec![4.0])),
            }),
        }),
    })
    .unwrap();
    assert_eq!(unbalanced.nodes.len(), 7);
    assert_eq!(
        unbalanced
            .nodes
            .iter()
            .map(|node| node.id)
            .collect::<Vec<_>>(),
        (0..7).collect::<Vec<_>>()
    );
}

#[test]
fn rejects_inconsistent_canonical_models() {
    assert!(LogisticModelWeights::new(array![[1.0, 2.0]], array![0.0, 1.0], 2).is_err());
    let tree = flatten_tree(&RecursiveNode::Leaf(vec![1.0])).unwrap();
    let forest = ForestStructure {
        trees: vec![tree],
        aggregation: AggregationMode::Sum,
        n_targets: 2,
    };
    assert!(export_tree_ensemble(&forest, TreeTask::Regression).is_err());
}

#[test]
fn exports_svm_regressor() {
    let svm = SvmRegressor::new(
        array![[1.0, 0.0], [0.0, 1.0]],
        array![0.75, -0.25],
        -0.5,
        SvmKernel::Rbf { gamma: 0.2 },
        false,
    )
    .unwrap();
    let graph = export_svm_regressor(&svm).graph.unwrap();
    assert_eq!(graph.node[0].op_type, "SVMRegressor");
    assert_eq!(graph.node[0].domain, "ai.onnx.ml");
    let kernel = graph.node[0]
        .attribute
        .iter()
        .find(|attribute| attribute.name == "kernel_type")
        .unwrap();
    assert_eq!(kernel.s, b"RBF");
}

#[test]
fn exports_compact_linear_svm_regressor() {
    let svm = SvmRegressor::new(
        ndarray::Array2::zeros((0, 2)),
        array![1.0, -2.0],
        0.5,
        SvmKernel::Linear,
        false,
    )
    .unwrap();
    let model = export_svm_regressor(&svm);
    let node = &model.graph.as_ref().unwrap().node[0];
    let n_supports = node
        .attribute
        .iter()
        .find(|attribute| attribute.name == "n_supports")
        .unwrap();
    assert_eq!(n_supports.i, 0);
    assert_eq!(
        model.graph.unwrap().input[0]
            .r#type
            .as_ref()
            .unwrap()
            .tensor_type
            .as_ref()
            .unwrap()
            .shape
            .as_ref()
            .unwrap()
            .dim[1]
            .value,
        Some(Value::DimValue(2))
    );
}

#[test]
fn rejects_invalid_svm_shapes() {
    assert!(SvmRegressor::new(
        ndarray::Array2::zeros((0, 2)),
        array![1.0],
        0.0,
        SvmKernel::Linear,
        false,
    )
    .is_err());
}

#[test]
fn exports_svm_classifier_with_labels_and_scores() {
    let svm = SvmClassifier {
        support_vectors: array![[-1.0, 0.0], [1.0, 0.0]],
        coefficients: array![0.5, -0.5],
        rho: array![0.0],
        vectors_per_class: vec![1, 1],
        class_labels: vec![10, 20],
        prob_a: Vec::new(),
        prob_b: Vec::new(),
        kernel: SvmKernel::Linear,
    };
    let graph = export_svm_classifier(&svm).unwrap().graph.unwrap();
    assert_eq!(graph.node[0].op_type, "SVMClassifier");
    assert_eq!(graph.output.len(), 2);
    assert_eq!(
        graph.output[0]
            .r#type
            .as_ref()
            .unwrap()
            .tensor_type
            .as_ref()
            .unwrap()
            .elem_type,
        onnx_export_rs::graph_builder::INT64
    );
}

#[test]
fn exports_gradient_boosting_with_scaled_leaves_and_base_score() {
    let tree = flatten_tree(&RecursiveNode::Branch {
        feature: 0,
        threshold: 0.0,
        left: Box::new(RecursiveNode::Leaf(vec![2.0])),
        right: Box::new(RecursiveNode::Leaf(vec![-1.0])),
    })
    .unwrap();
    let boosted = GradientBoostedEnsemble {
        trees: vec![tree],
        base_values: vec![0.25],
        learning_rate: 0.1,
        n_targets: 1,
        task: TreeTask::Regression,
        post_transform: PostTransform::None,
    };
    let graph = export_gradient_boosting(&boosted).unwrap().graph.unwrap();
    let attributes = &graph.node[0].attribute;
    let base = attributes
        .iter()
        .find(|attribute| attribute.name == "base_values")
        .unwrap();
    let weights = attributes
        .iter()
        .find(|attribute| attribute.name == "target_weights")
        .unwrap();
    assert_eq!(base.floats, [0.25]);
    assert_eq!(weights.floats, [0.2, -0.1]);
}

#[cfg(feature = "validate")]
#[test]
fn zero_and_wide_range_linear_models_round_trip() {
    let zero = LinearModelWeights::new(array![0.0], 0.0);
    let zero_report = onnx_export_rs::validate::validate_export(
        &to_bytes(&export_linear(&zero)).unwrap(),
        &array![[-10.0], [0.0], [10.0]],
        |_| array![0.0, 0.0, 0.0],
        0.0,
    )
    .unwrap();
    assert!(zero_report.passed);

    let wide = LinearModelWeights::new(array![1e-8, 1e8], -3.0);
    let inputs = array![[1e8, 1e-8], [-1e8, -1e-8]];
    let report = onnx_export_rs::validate::validate_export(
        &to_bytes(&export_linear(&wide)).unwrap(),
        &inputs,
        |values| values.dot(&array![1e-8, 1e8]) - 3.0,
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(feature = "validate")]
#[test]
fn validation_detects_corruption() {
    let report = onnx_export_rs::validate::compare_predictions(&[1.0, 2.0], &[1.0, 2.5], 1e-5);
    assert!(!report.passed);
    assert_eq!(report.max_absolute_difference, 0.5);
}

#[cfg(feature = "validate")]
#[test]
fn tract_runs_exported_linear_model() {
    let weights = LinearModelWeights::new(array![2.0, -1.0], 0.5);
    let bytes = to_bytes(&export_linear(&weights)).unwrap();
    let inputs = array![[1.0, 3.0], [2.0, -1.0]];
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        |values| values.dot(&array![2.0, -1.0]) + 0.5,
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(feature = "validate")]
#[test]
fn tract_runs_exported_binary_logistic_model() {
    let weights = LogisticModelWeights::new(array![[1.0, -2.0]], array![0.25], 2).unwrap();
    let bytes = to_bytes(&export_logistic(&weights)).unwrap();
    let inputs = array![[1.0, 0.0], [0.0, 1.0], [2.0, -1.0]];
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        |values| {
            values
                .dot(&array![1.0, -2.0])
                .mapv(|score| 1.0 / (1.0 + (-(score + 0.25_f64)).exp()))
        },
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(feature = "validate")]
#[test]
fn tract_reports_tree_execution_as_unsupported() {
    let tree = flatten_tree(&RecursiveNode::Branch {
        feature: 0,
        threshold: 1.5,
        left: Box::new(RecursiveNode::Leaf(vec![10.0])),
        right: Box::new(RecursiveNode::Leaf(vec![-3.0])),
    })
    .unwrap();
    let forest = ForestStructure {
        trees: vec![tree],
        aggregation: AggregationMode::Average,
        n_targets: 1,
    };
    let bytes = to_bytes(&export_tree_ensemble(&forest, TreeTask::Regression).unwrap()).unwrap();
    let inputs = array![[1.0], [2.0]];
    let error =
        onnx_export_rs::validate::validate_export(&bytes, &inputs, |_| array![10.0, -3.0], 1e-5)
            .unwrap_err();
    assert!(error.to_string().contains("TreeEnsembleRegressor"));
}

#[cfg(feature = "smartcore")]
#[test]
fn adapts_smartcore_linear_regression() {
    use onnx_export_rs::adapters::smartcore::linear_weights;
    use smartcore::linalg::basic::matrix::DenseMatrix;
    use smartcore::linear::linear_regression::{LinearRegression, LinearRegressionParameters};

    let input =
        DenseMatrix::from_2d_array(&[&[0.0, 0.0], &[1.0, 0.0], &[0.0, 1.0], &[1.0, 1.0]]).unwrap();
    let targets = vec![0.5, 2.5, -0.5, 1.5];
    let fitted =
        LinearRegression::fit(&input, &targets, LinearRegressionParameters::default()).unwrap();
    let weights = linear_weights(&fitted);
    assert!((weights.intercept - 0.5).abs() < 1e-10);
    assert!((weights.coefficients[0] - 2.0).abs() < 1e-10);
    assert!((weights.coefficients[1] + 1.0).abs() < 1e-10);
}

#[cfg(feature = "smartcore")]
#[test]
fn adapts_smartcore_logistic_regression() {
    use onnx_export_rs::adapters::smartcore::{logistic_classes, logistic_weights};
    use smartcore::linalg::basic::matrix::DenseMatrix;
    use smartcore::linear::logistic_regression::{
        LogisticRegression, LogisticRegressionParameters,
    };

    let input = DenseMatrix::from_2d_array(&[
        &[-2.0, -1.0],
        &[-1.0, -2.0],
        &[-1.0, 0.0],
        &[1.0, 0.0],
        &[1.0, 2.0],
        &[2.0, 1.0],
    ])
    .unwrap();
    let targets = vec![10, 10, 10, 20, 20, 20];
    let fitted =
        LogisticRegression::fit(&input, &targets, LogisticRegressionParameters::default()).unwrap();
    let weights = logistic_weights(&fitted).unwrap();
    assert_eq!(weights.n_classes, 2);
    assert_eq!(weights.coefficients.dim(), (1, 2));
    assert_eq!(logistic_classes(&fitted), [10, 20]);
}

#[cfg(all(feature = "smartcore", feature = "validate"))]
#[test]
fn smartcore_logistic_export_runs_in_tract() {
    use onnx_export_rs::adapters::smartcore::logistic_weights;
    use smartcore::linalg::basic::matrix::DenseMatrix;
    use smartcore::linear::logistic_regression::LogisticRegression;

    let rows = [
        &[-2.0, -1.0][..],
        &[-1.0, -2.0],
        &[-1.0, 0.0],
        &[1.0, 0.0],
        &[1.0, 2.0],
        &[2.0, 1.0],
    ];
    let input = DenseMatrix::from_2d_array(&rows).unwrap();
    let fitted =
        LogisticRegression::fit(&input, &vec![0, 0, 0, 1, 1, 1], Default::default()).unwrap();
    let weights = logistic_weights(&fitted).unwrap();
    let bytes = to_bytes(&export_logistic(&weights)).unwrap();
    let inputs = array![[-2.0, -1.0], [2.0, 1.0]];
    let coefficients = weights.coefficients.row(0).to_owned();
    let intercept = weights.intercept[0];
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        move |values| {
            values
                .dot(&coefficients)
                .mapv(|score| 1.0 / (1.0 + (-(score + intercept)).exp()))
        },
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(all(feature = "smartcore", feature = "validate"))]
#[test]
fn smartcore_multiclass_logistic_export_runs_in_tract() {
    use ndarray::Array1;
    use onnx_export_rs::adapters::smartcore::{logistic_classes, logistic_weights};
    use smartcore::linalg::basic::matrix::DenseMatrix;
    use smartcore::linear::logistic_regression::LogisticRegression;

    let training = [
        &[-3.0, 0.0][..],
        &[-2.0, -1.0],
        &[-2.0, 1.0],
        &[0.0, 3.0],
        &[-1.0, 2.0],
        &[1.0, 2.0],
        &[3.0, 0.0],
        &[2.0, -1.0],
        &[2.0, 1.0],
    ];
    let matrix = DenseMatrix::from_2d_array(&training).unwrap();
    let fitted = LogisticRegression::fit(
        &matrix,
        &vec![10, 10, 10, 20, 20, 20, 30, 30, 30],
        Default::default(),
    )
    .unwrap();
    let weights = logistic_weights(&fitted).unwrap();
    assert_eq!(logistic_classes(&fitted), [10, 20, 30]);
    let bytes = to_bytes(&export_logistic(&weights)).unwrap();
    let inputs = array![[-2.5, 0.0], [0.0, 2.5], [2.5, 0.0]];
    let coefficients = weights.coefficients.clone();
    let intercept = weights.intercept.clone();
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        move |values| {
            let scores = values.dot(&coefficients.t()) + &intercept;
            Array1::from_iter(scores.rows().into_iter().flat_map(|row| {
                let maximum = row.iter().copied().fold(f64::NEG_INFINITY, f64::max);
                let denominator = row.iter().map(|value| (value - maximum).exp()).sum::<f64>();
                row.iter()
                    .map(move |value| (value - maximum).exp() / denominator)
                    .collect::<Vec<_>>()
            }))
        },
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(feature = "smartcore")]
#[test]
fn adapts_smartcore_regularized_regressions() {
    use onnx_export_rs::adapters::smartcore::{elastic_net_weights, lasso_weights, ridge_weights};
    use smartcore::linalg::basic::matrix::DenseMatrix;
    use smartcore::linear::elastic_net::{ElasticNet, ElasticNetParameters};
    use smartcore::linear::lasso::{Lasso, LassoParameters};
    use smartcore::linear::ridge_regression::{RidgeRegression, RidgeRegressionParameters};

    let input = DenseMatrix::from_2d_array(&[
        &[0.0, 0.0],
        &[1.0, 0.0],
        &[0.0, 1.0],
        &[1.0, 1.0],
        &[2.0, 1.0],
        &[1.0, 2.0],
    ])
    .unwrap();
    let targets = vec![0.5, 2.5, -0.5, 1.5, 3.5, 0.5];
    let ridge =
        RidgeRegression::fit(&input, &targets, RidgeRegressionParameters::default()).unwrap();
    let lasso = Lasso::fit(&input, &targets, LassoParameters::default()).unwrap();
    let elastic = ElasticNet::fit(&input, &targets, ElasticNetParameters::default()).unwrap();
    let adapted = [
        (ridge_weights(&ridge), ridge.predict(&input).unwrap()),
        (lasso_weights(&lasso), lasso.predict(&input).unwrap()),
        (
            elastic_net_weights(&elastic),
            elastic.predict(&input).unwrap(),
        ),
    ];
    for (weights, predictions) in adapted {
        assert_eq!(weights.n_features(), 2);
        assert!(weights.intercept.is_finite());
        assert!(weights.coefficients.iter().all(|value| value.is_finite()));
        for (row, prediction) in [
            [0.0, 0.0],
            [1.0, 0.0],
            [0.0, 1.0],
            [1.0, 1.0],
            [2.0, 1.0],
            [1.0, 2.0],
        ]
        .iter()
        .zip(predictions)
        {
            let canonical = weights.intercept
                + row[0] * weights.coefficients[0]
                + row[1] * weights.coefficients[1];
            assert!((canonical - prediction).abs() < 1e-10);
        }
    }
}

#[cfg(all(feature = "smartcore", feature = "validate"))]
#[test]
fn adapts_and_runs_smartcore_pca_and_svd() {
    use onnx_export_rs::adapters::smartcore::{pca_transform, svd_transform};
    use onnx_export_rs::exporters::export_affine;
    use smartcore::decomposition::pca::{PCAParameters, PCA};
    use smartcore::decomposition::svd::{SVDParameters, SVD};
    use smartcore::linalg::basic::arrays::Array;
    use smartcore::linalg::basic::matrix::DenseMatrix;

    let source = DenseMatrix::from_2d_array(&[
        &[1.0, 2.0, 0.0],
        &[2.0, 1.0, 1.0],
        &[3.0, 4.0, 2.0],
        &[4.0, 3.0, 4.0],
    ])
    .unwrap();
    let inputs = array![[1.5, 2.5, 0.5], [3.5, 3.5, 3.0]];
    let probe = DenseMatrix::from_2d_array(&[&[1.5, 2.5, 0.5], &[3.5, 3.5, 3.0]]).unwrap();

    let pca = PCA::fit(&source, PCAParameters::default().with_n_components(2)).unwrap();
    let svd = SVD::fit(&source, SVDParameters::default().with_n_components(2)).unwrap();
    for (canonical, expected) in [
        (pca_transform(&pca).unwrap(), pca.transform(&probe).unwrap()),
        (svd_transform(&svd).unwrap(), svd.transform(&probe).unwrap()),
    ] {
        let expected = ndarray::Array1::from_iter(expected.iterator(0).copied());
        let report = onnx_export_rs::validate::validate_export(
            &to_bytes(&export_affine(&canonical)).unwrap(),
            &inputs,
            move |_| expected.clone(),
            1e-5,
        )
        .unwrap();
        assert!(report.passed, "{report:?}");
    }
}

#[cfg(feature = "linfa")]
#[test]
fn adapts_linfa_linear_regression() {
    use linfa::traits::Fit;
    use onnx_export_rs::adapters::linfa::linear_weights;

    let dataset = linfa::Dataset::new(
        array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
        array![0.5, 2.5, -0.5, 1.5],
    );
    let fitted = linfa_linear::LinearRegression::new().fit(&dataset).unwrap();
    let weights = linear_weights(&fitted);
    assert!((weights.intercept - 0.5).abs() < 1e-10);
    assert!((weights.coefficients[0] - 2.0).abs() < 1e-10);
    assert!((weights.coefficients[1] + 1.0).abs() < 1e-10);
}

#[cfg(all(feature = "linfa", feature = "validate"))]
#[test]
fn adapts_and_runs_linfa_binary_logistic_regression() {
    use linfa::traits::Fit;
    use onnx_export_rs::adapters::linfa::binary_logistic_weights;

    let dataset = linfa::Dataset::new(
        array![
            [-2.0, -1.0],
            [-1.0, -2.0],
            [-1.0, 0.0],
            [1.0, 0.0],
            [1.0, 2.0],
            [2.0, 1.0]
        ],
        array![10, 10, 10, 20, 20, 20],
    );
    let fitted = linfa_logistic::LogisticRegression::default()
        .fit(&dataset)
        .unwrap();
    let weights = binary_logistic_weights(&fitted).unwrap();
    let bytes = to_bytes(&export_logistic(&weights)).unwrap();
    let inputs = array![[-2.0, -1.0], [2.0, 1.0]];
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        |values| fitted.predict_probabilities(values),
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(all(feature = "linfa", feature = "validate"))]
#[test]
fn adapts_and_runs_linfa_multiclass_logistic_regression() {
    use linfa::traits::Fit;
    use ndarray::Array1;
    use onnx_export_rs::adapters::linfa::{
        multiclass_logistic_classes, multiclass_logistic_weights,
    };

    let dataset = linfa::Dataset::new(
        array![
            [-3.0, 0.0],
            [-2.0, -1.0],
            [-2.0, 1.0],
            [0.0, 3.0],
            [-1.0, 2.0],
            [1.0, 2.0],
            [3.0, 0.0],
            [2.0, -1.0],
            [2.0, 1.0]
        ],
        array![10, 10, 10, 20, 20, 20, 30, 30, 30],
    );
    let fitted = linfa_logistic::MultiLogisticRegression::default()
        .fit(&dataset)
        .unwrap();
    let weights = multiclass_logistic_weights(&fitted).unwrap();
    assert_eq!(multiclass_logistic_classes(&fitted), [10, 20, 30]);
    let bytes = to_bytes(&export_logistic(&weights)).unwrap();
    let inputs = array![[-2.5, 0.0], [0.0, 2.5], [2.5, 0.0]];
    let report = onnx_export_rs::validate::validate_export(
        &bytes,
        &inputs,
        |values| Array1::from_iter(fitted.predict_probabilities(values)),
        1e-5,
    )
    .unwrap();
    assert!(report.passed, "{report:?}");
}

#[cfg(feature = "linfa")]
#[test]
fn adapts_linfa_classification_tree() {
    use linfa::traits::{Fit, Predict};
    use onnx_export_rs::adapters::linfa::classification_tree;
    use onnx_export_rs::canonical::BranchMode;

    let records = array![
        [-3.0, 0.0],
        [-2.0, 1.0],
        [-1.0, -1.0],
        [1.0, -1.0],
        [2.0, 1.0],
        [3.0, 0.0]
    ];
    let dataset = linfa::Dataset::new(records.clone(), array![0, 0, 0, 1, 1, 1]);
    let fitted = linfa_trees::DecisionTree::params()
        .max_depth(Some(3))
        .fit(&dataset)
        .unwrap();
    let canonical = classification_tree(&fitted, &[0, 1]).unwrap();
    assert!(canonical
        .nodes
        .iter()
        .filter(|node| !node.is_leaf())
        .all(|node| node.branch_mode == BranchMode::LessThan));
    let forest = ForestStructure {
        trees: vec![canonical.clone()],
        aggregation: AggregationMode::Sum,
        n_targets: 2,
    };
    let model = export_tree_ensemble(&forest, TreeTask::Classification).unwrap();
    let graph = model.graph.unwrap();
    assert_eq!(graph.node[1].op_type, "ArgMax");
    assert_eq!(
        graph.output[0]
            .r#type
            .as_ref()
            .unwrap()
            .tensor_type
            .as_ref()
            .unwrap()
            .elem_type,
        onnx_export_rs::graph_builder::INT64
    );
    let modes = &graph.node[0]
        .attribute
        .iter()
        .find(|attribute| attribute.name == "nodes_modes")
        .unwrap()
        .strings;
    assert!(modes.iter().any(|mode| mode == b"BRANCH_LT"));

    let expected = fitted.predict(&records);
    for (row, expected_class) in records.rows().into_iter().zip(expected) {
        let mut node_id = 0_usize;
        loop {
            let node = &canonical.nodes[node_id];
            if node.is_leaf() {
                let predicted = node
                    .leaf_values
                    .iter()
                    .enumerate()
                    .max_by(|left, right| left.1.total_cmp(right.1))
                    .unwrap()
                    .0;
                assert_eq!(predicted, expected_class);
                break;
            }
            let feature_value = row[node.feature_id as usize];
            let take_true = match node.branch_mode {
                BranchMode::LessOrEqual => feature_value <= f64::from(node.threshold),
                BranchMode::LessThan => feature_value < f64::from(node.threshold),
            };
            node_id = if take_true {
                node.true_child_id as usize
            } else {
                node.false_child_id as usize
            };
        }
    }
}

#[cfg(feature = "linfa")]
#[test]
fn adapts_extended_linfa_models() {
    use linfa::traits::{Fit, FitWith};
    use onnx_export_rs::adapters::linfa::{
        elastic_net_weights, ftrl, gaussian_mixture, kmeans, linear_svm_score, pca,
    };

    let records = array![
        [-3.0, 0.0, 1.0],
        [-2.0, 1.0, 0.0],
        [2.0, 0.0, 1.0],
        [3.0, 1.0, 0.0]
    ];
    let unsupervised = linfa::DatasetBase::from(records.clone());
    let clusters = linfa_clustering::KMeans::params(2)
        .fit(&unsupervised)
        .unwrap();
    assert_eq!(kmeans(&clusters).unwrap().centroids.dim(), (2, 3));

    let mixture = linfa_clustering::GaussianMixtureModel::params(2)
        .fit(&unsupervised)
        .unwrap();
    assert_eq!(gaussian_mixture(&mixture).unwrap().means.dim(), (2, 3));

    let reduction = linfa_reduction::Pca::params(2).fit(&unsupervised).unwrap();
    let projection = pca(&reduction).unwrap();
    assert_eq!(projection.matrix.dim(), (3, 2));

    let regression = linfa::Dataset::new(records, array![-2.0, -1.0, 1.0, 2.0]);
    let elastic = linfa_elasticnet::ElasticNet::params()
        .penalty(0.1)
        .fit(&regression)
        .unwrap();
    assert_eq!(elastic_net_weights(&elastic).n_features(), 3);

    let binary = linfa::Dataset::new(
        array![[-2.0, 0.0], [-1.0, 0.0], [1.0, 0.0], [2.0, 0.0]],
        array![false, false, true, true],
    );
    let online = linfa_ftrl::Ftrl::params().fit_with(None, &binary).unwrap();
    assert_eq!(ftrl(&online).unwrap().coefficients.dim(), (1, 2));

    let svm = linfa_svm::Svm::<_, bool>::params()
        .linear_kernel()
        .fit(&binary)
        .unwrap();
    assert_eq!(linear_svm_score(&svm, 2).unwrap().n_features(), 2);
}

#[cfg(feature = "linfa-compat")]
#[test]
fn adapts_private_linfa_bayes_and_pls_state() {
    use linfa::traits::Fit;
    use onnx_export_rs::adapters::linfa_compat;

    let classification = linfa::Dataset::new(
        array![[-2.0, 0.0], [-1.0, 0.5], [1.0, 0.0], [2.0, 0.5]],
        array![10, 10, 20, 20],
    );
    let gaussian = linfa_bayes::GaussianNb::params()
        .fit(&classification)
        .unwrap();
    assert_eq!(
        linfa_compat::gaussian_naive_bayes(&gaussian)
            .unwrap()
            .means
            .dim(),
        (2, 2)
    );

    let counts = linfa::Dataset::new(
        array![[2.0, 1.0], [3.0, 1.0], [0.0, 2.0], [0.0, 3.0]],
        array![10, 10, 20, 20],
    );
    let multinomial = linfa_bayes::MultinomialNb::params().fit(&counts).unwrap();
    assert_eq!(
        linfa_compat::multinomial_naive_bayes(&multinomial)
            .unwrap()
            .coefficients
            .dim(),
        (2, 2)
    );

    let regression = linfa::Dataset::new(
        array![[1.0, 2.0], [2.0, 1.0], [3.0, 4.0], [4.0, 3.0]],
        array![[1.0, 0.0], [2.0, 1.0], [3.0, 2.0], [4.0, 3.0]],
    );
    let pls = linfa_pls::PlsRegression::params(1)
        .fit(&regression)
        .unwrap();
    assert_eq!(linfa_compat::pls(&pls).unwrap().matrix.dim(), (2, 2));
}

#[cfg(feature = "linfa-compat")]
#[test]
fn adapts_private_linfa_nonlinear_svm_regressor() {
    use linfa::traits::{Fit, Predict};
    use onnx_export_rs::adapters::linfa_compat;
    use onnx_export_rs::canonical::SvmKernel;

    let data = linfa::Dataset::new(
        array![[-2.0], [-1.0], [0.0], [1.0], [2.0]],
        array![-1.5, -0.8, 0.1, 0.9, 1.4],
    );
    let svm = linfa_svm::Svm::<f64, f64>::params()
        .gaussian_kernel(0.5)
        .c_svr(10.0, Some(0.1))
        .fit(&data)
        .unwrap();

    let model = linfa_compat::svm_regressor(&svm).unwrap();
    let SvmKernel::Rbf { gamma } = model.kernel else {
        panic!("expected an RBF kernel, got {:?}", model.kernel);
    };
    // Linfa's Gaussian width eps = 0.5 maps to ONNX gamma = 1 / eps.
    assert!((gamma - 2.0).abs() < 1e-12);
    assert_eq!(model.support_vectors.nrows(), model.coefficients.len());

    // The reconstructed ONNX-ML score must reproduce Linfa's own prediction.
    for point in [-1.5_f64, -0.3, 0.7, 1.8] {
        let query = array![point];
        let expected = svm.predict(query.clone());
        let reconstructed = model.rho
            + model
                .support_vectors
                .outer_iter()
                .zip(&model.coefficients)
                .map(|(support, &coefficient)| {
                    let squared_distance =
                        (&support.to_owned() - &query).mapv(|value| value * value).sum();
                    coefficient * (-gamma * squared_distance).exp()
                })
                .sum::<f64>();
        assert!(
            (reconstructed - expected).abs() < 1e-9,
            "point {point}: {reconstructed} vs {expected}"
        );
    }

    // The polynomial kernel (<x, y> + c)^d carries an implicit unit gamma.
    let polynomial = linfa_svm::Svm::<f64, f64>::params()
        .polynomial_kernel(1.0, 2.0)
        .c_svr(10.0, Some(0.1))
        .fit(&data)
        .unwrap();
    let polynomial_model = linfa_compat::svm_regressor(&polynomial).unwrap();
    let SvmKernel::Polynomial {
        gamma: poly_gamma,
        coef0,
        degree,
    } = polynomial_model.kernel
    else {
        panic!("expected a polynomial kernel, got {:?}", polynomial_model.kernel);
    };
    assert_eq!((poly_gamma, coef0, degree), (1.0, 1.0, 2));
    for point in [-1.2_f64, 0.4, 1.6] {
        let query = array![point];
        let expected = polynomial.predict(query.clone());
        let reconstructed = polynomial_model.rho
            + polynomial_model
                .support_vectors
                .outer_iter()
                .zip(&polynomial_model.coefficients)
                .map(|(support, &coefficient)| {
                    let dot = (&support.to_owned() * &query).sum();
                    coefficient * (poly_gamma * dot + coef0).powi(degree as i32)
                })
                .sum::<f64>();
        assert!(
            (reconstructed - expected).abs() < 1e-9,
            "polynomial point {point}: {reconstructed} vs {expected}"
        );
    }
}

#[cfg(feature = "linfa-compat")]
#[test]
fn adapts_private_linfa_tweedie_glm_state() {
    use linfa::traits::Fit;
    use onnx_export_rs::adapters::linfa_compat;

    let counts = linfa::Dataset::new(
        array![[0.0, 1.0], [1.0, 0.0], [2.0, 1.0], [3.0, 2.0]],
        array![1.0, 2.0, 5.0, 12.0],
    );
    let fitted = linfa_linear::TweedieRegressor::params()
        .power(1.0)
        .link(linfa_linear::Link::Log)
        .alpha(0.0)
        .fit(&counts)
        .unwrap();

    let model = linfa_compat::tweedie_regressor(&fitted).unwrap();
    assert_eq!(model.link, LinkFunction::Log);
    // The serde round-trip must recover the fitted parameters exactly.
    assert_eq!(model.coefficients.as_slice().unwrap(), fitted.coef.as_slice().unwrap());
    assert_eq!(model.intercept, fitted.intercept);

    let nodes = export_generalized_linear(&model).graph.unwrap().node;
    assert_eq!(nodes[0].op_type, "Gemm");
    assert_eq!(nodes[1].op_type, "Exp");
}