tract-onnx 0.23.0-dev.2

use crate::model::{OnnxOpRegister, ParsingContext};
use crate::pb::*;
use tract_hir::internal::*;
use tract_onnx_opl::ml::*;

pub fn register_all_ops(reg: &mut OnnxOpRegister) {
    reg.insert("CategoryMapper", category_mapper);
}

#[derive(Debug, Clone, Hash)]
struct CategoryMapper {
    pub from: Arc<Tensor>,
    pub to: Arc<Tensor>,
    pub fallback: Arc<Tensor>,
}

impl Expansion for CategoryMapper {
    fn name(&self) -> StaticName {
        "CategoryMapper".into()
    }

    fn rules<'r, 'p: 'r, 's: 'r>(
        &'s self,
        s: &mut Solver<'r>,
        inputs: &'p [TensorProxy],
        outputs: &'p [TensorProxy],
    ) -> InferenceResult {
        check_input_arity(inputs, 1)?;
        check_output_arity(outputs, 1)?;
        s.equals(&inputs[0].shape, &outputs[0].shape)?;
        s.equals(&inputs[0].datum_type, self.from.datum_type())?;
        s.equals(&outputs[0].datum_type, self.to.datum_type())?;
        Ok(())
    }

    fn wire(
        &self,
        prefix: &str,
        model: &mut TypedModel,
        inputs: &[OutletId],
    ) -> TractResult<TVec<OutletId>> {
        let wire = model.wire_node(
            format!("{prefix}.reverse"),
            ReverseLookup::new(self.from.clone(), -1)?,
            inputs,
        )?;
        model.wire_node(
            format!("{prefix}.direct"),
            DirectLookup::new(self.to.clone(), self.fallback.clone())?,
            &wire,
        )
    }
}

fn category_mapper(
    _ctx: &ParsingContext,
    node: &NodeProto,
) -> TractResult<(Box<dyn InferenceOp>, Vec<String>)> {
    let ints: Vec<i64> = node.get_attr_vec("cats_int64s")?;
    let strings: Vec<String> = node.get_attr_vec("cats_strings")?;
    let default_int: Option<i64> = node.get_attr_opt("default_int64")?;
    let default_string: Option<String> = node.get_attr_opt("default_string")?;
    let op: Box<dyn InferenceOp> = match (default_int, default_string.as_ref()) {
        (None, None) | (Some(_), Some(_)) => bail!(
            "CategoryMapper requires exactly one of default_int64 and default_string (found {:?})",
            (default_int, default_string)
        ),
        (Some(def), None) => expand(CategoryMapper {
            from: rctensor1(&strings),
            to: rctensor1(&ints),
            fallback: rctensor0(def),
        }),
        (None, Some(def)) => expand(CategoryMapper {
            from: rctensor1(&ints),
            to: rctensor1(&strings),
            fallback: rctensor0(def.clone()),
        }),
    };
    Ok((op, vec![]))
}