onnx-export-rs 0.1.1

//! Low-level ONNX graph construction and serialization.

use std::path::Path;

use ndarray::ArrayD;
use prost::Message;

use crate::proto::{
    tensor_shape_proto, type_proto, AttributeProto, GraphProto, ModelProto, NodeProto,
    OperatorSetIdProto, TensorProto, TensorShapeProto, TypeProto, ValueInfoProto,
};
use crate::{Error, Result};

/// ONNX enum value for a 32-bit floating-point tensor.
pub const FLOAT: i32 = 1;
/// ONNX enum value for a signed 64-bit integer tensor.
pub const INT64: i32 = 7;

/// A shape dimension: fixed or symbolic.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Dimension {
    /// A fixed dimension.
    Fixed(usize),
    /// A symbolic dimension such as `batch`.
    Symbolic(String),
}

/// Creates an `f32` initializer tensor from an ndarray.
#[must_use]
pub fn make_tensor(name: impl Into<String>, data: &ArrayD<f32>) -> TensorProto {
    TensorProto {
        dims: data.shape().iter().map(|&d| d as i64).collect(),
        data_type: FLOAT,
        float_data: data.iter().copied().collect(),
        int32_data: Vec::new(),
        int64_data: Vec::new(),
        name: name.into(),
        raw_data: Vec::new(),
    }
}

/// Creates an `i64` initializer tensor with the supplied shape.
///
/// # Panics
///
/// Panics when `shape` does not contain exactly `data.len()` elements.
#[must_use]
pub fn make_i64_tensor(name: impl Into<String>, shape: &[usize], data: Vec<i64>) -> TensorProto {
    assert_eq!(shape.iter().product::<usize>(), data.len());
    TensorProto {
        dims: shape.iter().map(|&dimension| dimension as i64).collect(),
        data_type: INT64,
        float_data: Vec::new(),
        int32_data: Vec::new(),
        int64_data: data,
        name: name.into(),
        raw_data: Vec::new(),
    }
}

/// Creates tensor value information. Use a symbolic first dimension for a
/// model that accepts dynamic batch sizes.
#[must_use]
pub fn make_value_info(name: impl Into<String>, shape: &[Dimension]) -> ValueInfoProto {
    make_typed_value_info(name, shape, FLOAT)
}

/// Creates tensor value information with an explicit ONNX element type.
#[must_use]
pub fn make_typed_value_info(
    name: impl Into<String>,
    shape: &[Dimension],
    element_type: i32,
) -> ValueInfoProto {
    let dim = shape
        .iter()
        .map(|value| tensor_shape_proto::Dimension {
            value: Some(match value {
                Dimension::Fixed(value) => {
                    tensor_shape_proto::dimension::Value::DimValue(*value as i64)
                }
                Dimension::Symbolic(value) => {
                    tensor_shape_proto::dimension::Value::DimParam(value.clone())
                }
            }),
        })
        .collect();
    ValueInfoProto {
        name: name.into(),
        r#type: Some(TypeProto {
            tensor_type: Some(type_proto::Tensor {
                elem_type: element_type,
                shape: Some(TensorShapeProto { dim }),
            }),
        }),
        doc_string: String::new(),
    }
}

/// Creates a generic ONNX node in the core domain.
#[must_use]
pub fn make_node(
    op_type: impl Into<String>,
    inputs: impl IntoIterator<Item = impl Into<String>>,
    outputs: impl IntoIterator<Item = impl Into<String>>,
    attributes: Vec<AttributeProto>,
) -> NodeProto {
    NodeProto {
        input: inputs.into_iter().map(Into::into).collect(),
        output: outputs.into_iter().map(Into::into).collect(),
        name: String::new(),
        op_type: op_type.into(),
        attribute: attributes,
        doc_string: String::new(),
        domain: String::new(),
    }
}

/// Creates an integer attribute.
#[must_use]
pub fn int_attribute(name: impl Into<String>, value: i64) -> AttributeProto {
    AttributeProto {
        name: name.into(),
        i: value,
        r#type: 2,
        ..Default::default()
    }
}

/// Creates an integer-list attribute.
#[must_use]
pub fn ints_attribute(name: impl Into<String>, values: Vec<i64>) -> AttributeProto {
    AttributeProto {
        name: name.into(),
        ints: values,
        r#type: 7,
        ..Default::default()
    }
}

/// Creates a float-list attribute.
#[must_use]
pub fn floats_attribute(name: impl Into<String>, values: Vec<f32>) -> AttributeProto {
    AttributeProto {
        name: name.into(),
        floats: values,
        r#type: 6,
        ..Default::default()
    }
}

/// Creates a byte-string attribute.
#[must_use]
pub fn string_attribute(name: impl Into<String>, value: impl Into<Vec<u8>>) -> AttributeProto {
    AttributeProto {
        name: name.into(),
        s: value.into(),
        r#type: 3,
        ..Default::default()
    }
}

/// Creates a byte-string-list attribute.
#[must_use]
pub fn strings_attribute(name: impl Into<String>, values: Vec<Vec<u8>>) -> AttributeProto {
    AttributeProto {
        name: name.into(),
        strings: values,
        r#type: 8,
        ..Default::default()
    }
}

/// Wraps a graph in an ONNX model with a core operator-set import.
#[must_use]
pub fn assemble_model(graph: GraphProto, opset_version: i64, ir_version: i64) -> ModelProto {
    ModelProto {
        ir_version,
        producer_name: env!("CARGO_PKG_NAME").into(),
        producer_version: env!("CARGO_PKG_VERSION").into(),
        domain: String::new(),
        model_version: 1,
        doc_string: String::new(),
        graph: Some(graph),
        opset_import: vec![OperatorSetIdProto {
            domain: String::new(),
            version: opset_version,
        }],
    }
}

/// Serializes a model to ONNX protobuf bytes.
pub fn to_bytes(model: &ModelProto) -> Result<Vec<u8>> {
    let mut bytes = Vec::with_capacity(model.encoded_len());
    model.encode(&mut bytes)?;
    Ok(bytes)
}

/// Saves an ONNX model to a file.
pub fn save_to_file(model: &ModelProto, path: impl AsRef<Path>) -> Result<()> {
    let path = path.as_ref();
    std::fs::write(path, to_bytes(model)?).map_err(|source| Error::Io {
        path: path.to_path_buf(),
        source,
    })
}