onnx-runtime-shape-inference 0.1.0-dev.6

//! Schema-aware `Einsum` shape inference through the shared canonical planner.
//!
//! Parsing, validation, ellipsis expansion, diagonals, reductions, and output
//! ordering live in `onnx-runtime-ir::EinsumPlan`. This handler only supplies
//! shape-inference dimensions and resolves ellipsis broadcasts through
//! [`InferenceContext::broadcast_dim`], preserving the context's symbolic
//! representative and lineage bookkeeping.

use onnx_runtime_ir::{Attribute, EinsumInput, EinsumPlan, EinsumResolveError};

use crate::context::{InferenceContext, TypeInfo};
use crate::dim_expr::DimExpr;
use crate::error::ShapeInferError;
use crate::registry::InferenceRegistry;

fn einsum(ctx: &mut InferenceContext) -> Result<(), ShapeInferError> {
    let equation_bytes = match ctx.node.attr("equation") {
        Some(Attribute::String(bytes)) => bytes,
        _ => {
            return Err(ShapeInferError::MissingAttribute {
                op: ctx.op().to_owned(),
                attr: "equation".into(),
            });
        }
    };
    let equation = std::str::from_utf8(equation_bytes).map_err(|error| {
        let location = error.valid_up_to();
        let detail = match error.error_len() {
            Some(length) => format!(
                "attribute `equation` is not valid UTF-8: invalid byte sequence of length {length} starts at byte offset {location}"
            ),
            None => format!(
                "attribute `equation` is not valid UTF-8: incomplete byte sequence starts at byte offset {location}"
            ),
        };
        ShapeInferError::Invalid {
            op: ctx.op().to_owned(),
            detail,
        }
    })?;
    if ctx.num_outputs() != 1 {
        return Err(ShapeInferError::Invalid {
            op: ctx.op().to_owned(),
            detail: format!(
                "equation `{equation}` requires exactly 1 output, but the node declares {} outputs",
                ctx.num_outputs()
            ),
        });
    }

    // Clone before planning so output resolution can mutably use the context's
    // broadcast chokepoint without retaining immutable borrows into `ctx`.
    let input_types: Vec<Option<TypeInfo>> = (0..ctx.num_inputs())
        .map(|i| ctx.input_type(i).cloned())
        .collect();
    let inputs: Vec<_> = input_types
        .iter()
        .map(|input| {
            EinsumInput::from_optional(
                input.as_ref().map(|type_info| type_info.dtype),
                input.as_ref().map(|type_info| type_info.shape.as_slice()),
            )
        })
        .collect();
    let imported_opset = ctx.opset("");
    let plan = match EinsumPlan::build_for_opset(equation, &inputs, imported_opset) {
        Ok(plan) => plan,
        Err(error) if error.is_incomplete_metadata() => return Ok(()),
        Err(error) => {
            return Err(ShapeInferError::Invalid {
                op: ctx.op().to_owned(),
                detail: error.to_string(),
            });
        }
    };

    let input_shapes: Vec<&[DimExpr]> = input_types
        .iter()
        .enumerate()
        .map(|(input_index, input)| {
            input
                .as_ref()
                .ok_or_else(|| ShapeInferError::Invalid {
                    op: ctx.op().to_owned(),
                    detail: format!(
                        "shared Einsum plan admitted input #{input_index} without resolved type metadata"
                    ),
                })
                .map(|type_info| type_info.shape.as_slice())
        })
        .collect::<Result<_, _>>()?;
    let output_dtype = plan.dtype();
    let output_shape = match plan
        .resolve_output_shape(&input_shapes, |left, right| ctx.broadcast_dim(left, right))
    {
        Ok(shape) => shape,
        Err(EinsumResolveError::Broadcast { source, .. }) => return Err(source),
        Err(EinsumResolveError::InputCount { expected, found }) => {
            return Err(ShapeInferError::Invalid {
                op: ctx.op().to_owned(),
                detail: format!(
                    "shared Einsum plan expected {expected} resolved input shapes, found {found}"
                ),
            });
        }
        Err(EinsumResolveError::InputRank {
            input,
            expected,
            found,
        }) => {
            return Err(ShapeInferError::Invalid {
                op: ctx.op().to_owned(),
                detail: format!(
                    "shared Einsum plan expected input #{input} rank {expected}, found {found}"
                ),
            });
        }
    };

    ctx.set_output(0, output_dtype, output_shape);
    Ok(())
}

/// Register the `Einsum` rule.
pub fn register(reg: &mut InferenceRegistry) {
    reg.register("", "Einsum", 12, einsum);
}