rten 0.25.0

//! The [`Operator`] trait for defining operators.

use std::any::Any;
use std::borrow::Cow;
use std::convert::Infallible;
use std::error::Error;
use std::fmt;
use std::fmt::{Debug, Display};

use rten_base::bit_set::BitSet;
use rten_gemm::PackedBMatrix;
use rten_tensor::errors::DimensionError;
use rten_tensor::{Layout, Storage, TensorBase};
use smallvec::SmallVec;

use crate::BufferPool;
use crate::graph::{CaptureEnv, Graph, RunError, RunOptions};
use crate::infer_shapes::InferShapes;
use crate::timing::Profiler;
use crate::value::{DataType, DataTypeOf, TryFromValueError, Value, ValueType, ValueView};
use crate::weight_cache::WeightCache;

/// An operator input which has been pre-packed for more efficient use during
/// inference.
pub enum PrepackedInput {
    /// Prepacked RHS / B input for matrix multiplication with f32 weights.
    FloatBMatrix(PackedBMatrix<f32>),

    /// Prepacked RHS / B input for matrix multiplication with i8 weights.
    Int8BMatrix(PackedBMatrix<i8>),
}

impl PrepackedInput {
    fn dtype(&self) -> DataType {
        match self {
            Self::FloatBMatrix(_) => DataType::Float,
            Self::Int8BMatrix(_) => DataType::Int8,
        }
    }
}

macro_rules! impl_prepacked_input_conversions {
    ($type:ty, $variant:ident) => {
        impl From<PackedBMatrix<$type>> for PrepackedInput {
            fn from(value: PackedBMatrix<$type>) -> Self {
                PrepackedInput::$variant(value)
            }
        }

        impl<'a> TryFrom<&'a PrepackedInput> for &'a PackedBMatrix<$type> {
            type Error = TryFromValueError;

            fn try_from(ppi: &'a PrepackedInput) -> Result<Self, Self::Error> {
                match ppi {
                    PrepackedInput::$variant(packed) => Ok(packed),
                    _ => Err(TryFromValueError::WrongType {
                        actual: ValueType::Tensor(ppi.dtype()),
                        expected: ValueType::Tensor(<$type as DataTypeOf>::dtype_of()),
                    }),
                }
            }
        }
    };
}
impl_prepacked_input_conversions!(f32, FloatBMatrix);
impl_prepacked_input_conversions!(i8, Int8BMatrix);

/// Trait for values that can be converted into the result type used by
/// [`Operator::run`].
pub trait IntoOpResult {
    fn into_op_result(self) -> Result<OutputList, OpError>;
}

impl IntoOpResult for Result<Value, OpError> {
    fn into_op_result(self) -> Result<OutputList, OpError> {
        self.map(|out| [out].into())
    }
}

impl IntoOpResult for Value {
    fn into_op_result(self) -> Result<OutputList, OpError> {
        Ok([self].into())
    }
}

impl<S: Storage, L: Layout> IntoOpResult for TensorBase<S, L>
where
    Value: From<TensorBase<S, L>>,
{
    fn into_op_result(self) -> Result<OutputList, OpError> {
        let output: Value = self.into();
        Ok([output].into())
    }
}

impl<S: Storage, L: Layout> IntoOpResult for Result<TensorBase<S, L>, OpError>
where
    Value: From<TensorBase<S, L>>,
{
    fn into_op_result(self) -> Result<OutputList, OpError> {
        self.map(|tensor| [tensor.into()].into())
    }
}

impl<T> IntoOpResult for Result<Vec<T>, OpError>
where
    Value: From<T>,
{
    fn into_op_result(self) -> Result<OutputList, OpError> {
        self.map(|tensors| tensors.into_iter().map(|t| t.into()).collect())
    }
}

/// Possible reasons why an operator may fail on a given input.
#[derive(Eq, PartialEq, Debug)]
pub enum OpError {
    /// Casting a tensor to an expected type or rank failed.
    CastFailed(TryFromValueError),

    /// Casting an input to an expected type or rank failed.
    InputCastFailed {
        index: usize,
        error: TryFromValueError,
    },

    /// A tensor has an unsupported type.
    UnsupportedType,

    /// Input tensor shapes are not compatible with each other or operator
    /// attributes.
    IncompatibleInputShapes(&'static str),

    /// The number of inputs was less than the required number.
    MissingInputs,

    /// An input has a value that is incorrect.
    InvalidValue(&'static str),

    /// An input or attribute has a value that is valid, but not currently supported.
    UnsupportedValue(&'static str),
}

impl OpError {
    /// Associate this error with a given operator input.
    pub fn with_input_index(self, index: usize) -> OpError {
        match self {
            Self::CastFailed(error) => OpError::InputCastFailed { index, error },
            Self::InputCastFailed { error, .. } => OpError::InputCastFailed { index, error },
            other => other,
        }
    }
}

impl From<DimensionError> for OpError {
    fn from(val: DimensionError) -> OpError {
        OpError::CastFailed(val.into())
    }
}

impl From<TryFromValueError> for OpError {
    fn from(val: TryFromValueError) -> OpError {
        OpError::CastFailed(val)
    }
}

impl From<Infallible> for OpError {
    fn from(x: Infallible) -> OpError {
        match x {}
    }
}

impl Display for OpError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            OpError::CastFailed(err) => write!(f, "{}", err),
            OpError::InputCastFailed { index, error } => {
                write!(f, "conversion error for input {}: {}", index, error)
            }
            OpError::IncompatibleInputShapes(details) => {
                write!(f, "incompatible input shapes: {}", details)
            }
            OpError::MissingInputs => write!(f, "required inputs were missing"),
            OpError::InvalidValue(details) => {
                write!(f, "input or attribute has invalid value: {}", details)
            }
            OpError::UnsupportedValue(details) => {
                write!(f, "unsupported input or attribute value: {}", details)
            }
            OpError::UnsupportedType => {
                write!(f, "unsupported input type")
            }
        }
    }
}

impl Error for OpError {}

/// Convert a tensor with dynamic dimension count to a view with a static
/// dimension count.
///
/// If the conversion fails an `OpError::InvalidValue` error will be returned
/// with a message that includes the name of the tensor and, optionally, the
/// names of the expected dimensions (eg. "NCHW").
macro_rules! static_dims {
    ($tensor:ident, $ndim:literal, $dim_names:literal) => {{
        use rten_tensor::prelude::*;

        if $tensor.ndim() != $ndim {
            Err(OpError::InvalidValue(concat!(
                stringify!($tensor),
                " must have ",
                stringify!($ndim),
                " dims (",
                $dim_names,
                ")"
            )))
        } else {
            Ok($tensor.nd_view::<$ndim>())
        }
    }};

    ($tensor:ident, $ndim:literal) => {{
        use rten_tensor::prelude::*;

        if $tensor.ndim() != $ndim {
            Err(OpError::InvalidValue(concat!(
                stringify!($tensor),
                " must have ",
                stringify!($ndim),
                " dims"
            )))
        } else {
            Ok($tensor.nd_view::<$ndim>())
        }
    }};

    ($tensor:ident?, $ndim: expr) => {
        if let Some($tensor) = $tensor.as_ref() {
            Some(static_dims!($tensor, $ndim))
        } else {
            None
        }
    };
}

pub(crate) use static_dims;

/// Check that two input dimension sizes are equal.
macro_rules! check_eq {
    ($actual:expr, $expected:expr) => {{
        if $actual != $expected {
            return Err(OpError::IncompatibleInputShapes(stringify!(
                $actual != $expected
            )));
        }
        Ok::<_, OpError>(())
    }};
}
pub(crate) use check_eq;

/// The number of outputs an operator has and which are used.
///
/// This contains a mask that tracks which of the first 32 outputs are used,
/// plus the total number. Outputs beyond the first 32 are always treated as
/// used.
#[derive(Copy, Clone, Debug, Default, PartialEq)]
pub struct OutputMask {
    /// Mask tracking which of the first 32 outputs are used.
    mask: BitSet<u32>,
    /// Total number of outputs.
    len: u32,
}

impl OutputMask {
    /// Return a mask with `len` outputs, where `mask` specifies which of the
    /// first 32 are used.
    pub fn new(mask: BitSet<u32>, len: u32) -> Self {
        Self { mask, len }
    }

    /// Return a mask with `len` outputs, all used.
    pub fn all_used(len: usize) -> Self {
        Self {
            mask: BitSet::ones(len.min(u32::BITS as usize) as u32),
            len: len as u32,
        }
    }

    /// Return the total number of outputs.
    pub fn len(&self) -> usize {
        self.len as usize
    }

    /// Return true if the output at position `idx` is used.
    ///
    /// Returns false if `idx` is out of bounds.
    pub fn is_used(&self, idx: usize) -> bool {
        if idx < u32::BITS as usize {
            self.mask.get(idx as u32)
        } else {
            idx < self.len as usize
        }
    }
}

/// Context passed to [`Operator::run`] containing the information needed for
/// the operator to execute.
pub struct OpRunContext<'a, 'i> {
    pool: &'a BufferPool,
    inputs: &'a InputList<'i>,
    outputs: OutputMask,
    name: Option<&'a str>,
}

impl<'a, 'i> OpRunContext<'a, 'i> {
    /// Create a new context.
    ///
    /// `outputs` is a mask indicating which of the operator's outputs are
    /// requested.
    pub fn new(pool: &'a BufferPool, inputs: &'a InputList<'i>, outputs: OutputMask) -> Self {
        OpRunContext {
            pool,
            inputs,
            outputs,
            name: None,
        }
    }

    /// Construct a new context with the same properties but different inputs.
    ///
    /// This is useful when one operator wants to delegate to another.
    pub fn with_new_inputs<'b, 'il>(&self, inputs: &'b InputList<'il>) -> OpRunContext<'b, 'il>
    where
        'a: 'b,
    {
        OpRunContext { inputs, ..*self }
    }

    /// The pool which should be used to allocate large buffers.
    pub fn pool(&self) -> &BufferPool {
        self.pool
    }

    /// Inputs to the operator execution.
    ///
    /// For in-place execution via [`Operator::run_in_place`] this contains
    /// the non in-place inputs.
    pub fn inputs(&self) -> &InputList<'i> {
        self.inputs
    }

    /// Return a mask indicating the requested outputs.
    ///
    /// This can be used to skip generating outputs that are unused, or in
    /// the rare cases that the output count cannot be determined from the
    /// operator's inputs and attributes alone.
    pub fn outputs(&self) -> OutputMask {
        self.outputs
    }

    /// Set the name of the current node in the graph.
    pub fn set_name(&mut self, name: Option<&'a str>) {
        self.name = name;
    }

    /// Return the name of the current node in the graph.
    pub fn name(&self) -> Option<&str> {
        self.name
    }
}

/// Rule for determining the type of an operator output.
#[derive(Copy, Clone)]
pub enum OutputType {
    /// This output has a fixed type, given the operator's attributes.
    Fixed(ValueType),
    /// This output has the same type as the input at a given index.
    CopyFromInput(u32),
    /// The output is the element type of an input sequence.
    ElementTypeOfInputSequence(u32),
    /// The output is a sequence whose element type matches an input.
    SequenceWithElementTypeOfInput(u32),
}

/// List of type rules for each operator output.
pub type OutputTypeList = SmallVec<[OutputType; 1]>;

/// Context passed to [`Operator::output_types`].
pub struct OutputTypesContext {
    /// Number of output value nodes connected to this operator.
    pub num_outputs: usize,
}

/// Outputs from an operator.
///
/// This avoids allocations in the common case where an operator produces
/// exactly one output.
pub type OutputList = SmallVec<[Value; 1]>;

/// Owned input values which an operator should modify in-place.
///
/// This is a list of `(input_index, value)` for owned inputs passed to
/// [`Operator::run_in_place`].
#[derive(Debug)]
pub struct InPlaceInputs {
    inputs: SmallVec<[(usize, Value); 1]>,
}

impl InPlaceInputs {
    /// Return the number of in-place inputs.
    pub fn len(&self) -> usize {
        self.inputs.len()
    }

    /// Return true if there are no in-place inputs.
    pub fn is_empty(&self) -> bool {
        self.inputs.is_empty()
    }

    /// Return the single in-place input.
    ///
    /// Panics if the list does not have exactly one element.
    pub fn into_single(self) -> Value {
        assert_eq!(
            self.inputs.len(),
            1,
            "expected exactly one in-place input, got {}",
            self.inputs.len()
        );
        self.inputs.into_iter().next().unwrap().1
    }
}

impl From<(usize, Value)> for InPlaceInputs {
    fn from(input: (usize, Value)) -> Self {
        InPlaceInputs {
            inputs: [input].into_iter().collect(),
        }
    }
}

impl FromIterator<(usize, Value)> for InPlaceInputs {
    fn from_iter<I: IntoIterator<Item = (usize, Value)>>(iter: I) -> Self {
        InPlaceInputs {
            inputs: iter.into_iter().collect(),
        }
    }
}

impl IntoIterator for InPlaceInputs {
    type Item = (usize, Value);
    type IntoIter = smallvec::IntoIter<[(usize, Value); 1]>;

    fn into_iter(self) -> Self::IntoIter {
        self.inputs.into_iter()
    }
}

/// An Operator performs a computation step when executing a data flow graph.
///
/// Operators take zero or more dynamic input values, plus a set of static
/// attributes and produce one or more output values.
///
/// Operators are usually named after the ONNX operator that they implement.
/// See <https://onnx.ai/onnx/operators/>.
pub trait Operator: Any + Debug {
    /// Return a display name for the operator.
    fn name(&self) -> &str;

    /// Execute the operator.
    ///
    /// `ctx` provides access to operator inputs and the [`BufferPool`] from
    /// which the output and temporary buffers should be allocated.
    ///
    /// For operators which have subgraphs (see
    /// [`as_subgraph_op`](Operator::as_subgraph_op)), the
    /// [`SubgraphOperator::run_subgraph`] method should be used instead.
    fn run(&self, ctx: &OpRunContext) -> Result<OutputList, OpError>;

    /// Return the maximum number of inputs this operator accepts.
    ///
    /// This can return `None` for variadic inputs with no limit.
    fn max_inputs(&self) -> Option<usize>;

    /// Return the maximum number of outputs this operator can produce.
    ///
    /// The default is one. Operators can return `None` to indicate no limit
    /// (eg. for variadic operators).
    fn max_outputs(&self) -> Option<usize> {
        Some(1)
    }

    /// Return the rules for determining the types of this operator's outputs.
    fn output_types(&self, ctx: &OutputTypesContext) -> Option<OutputTypeList>;

    /// Return the set of inputs which this operator can modify in-place via
    /// [`run_in_place`](Operator::run_in_place).
    ///
    /// In-place execution returns results by modifying an existing input
    /// instead of allocating a new tensor. Reducing memory allocations can
    /// significantly speed up graph runs.
    ///
    /// The returned bit set contains the index of each input the operator is
    /// able to modify. An empty set (the default) means in-place execution is
    /// not supported. Only the first 16 inputs are candidates for in-place
    /// modification.
    fn in_place_inputs(&self) -> BitSet<u16> {
        BitSet::new()
    }

    /// Return true if this operator is commutative, meaning that its inputs
    /// can be re-ordered without affecting the result.
    ///
    /// If true, the graph executor may swap inputs before calling the
    /// [`Operator::run_in_place`] implementation.
    fn is_commutative(&self) -> bool {
        false
    }

    /// Return true if this operator is associative, meaning that the operands
    /// of nested chains of this operator can be regrouped without affecting the
    /// result. Formally, `Op(Op(a, b), c) == Op(a, Op(b, c))`.
    ///
    /// The term _associative_ here refers to the ideal mathematical behavior of
    /// the operation, ignoring the non-associativity of corresponding
    /// floating-point operations.
    fn is_associative(&self) -> bool {
        false
    }

    /// Return true if this operator's outputs depend only on its inputs.
    ///
    /// The default implementation returns true, since most operators are
    /// deterministic. Operators such as random number generators however are
    /// not.
    ///
    /// The definition of _deterministic_ used here excludes minor differences
    /// due to eg. the order in which results from parallel sub-problems are
    /// accumulated. It also does not guarantee exact consistency across devices.
    fn is_deterministic(&self) -> bool {
        true
    }

    /// Execute this operator in-place on one or more existing tensors.
    ///
    /// This may only be called if `in_place_inputs` returns a non-empty set.
    ///
    /// `in_place` holds the owned values for the inputs from that set, which the
    /// implementation may modify and return as outputs. `ctx.inputs()` contains
    /// all of the operator's inputs, with the positions of the in-place inputs
    /// set to `None`.
    ///
    /// Operators may fall back to allocating a new output if some property of
    /// the input data or shapes means in-place operation is not possible. In
    /// this case they should return the input buffer to the pool, and allocate
    /// the new output buffer from it. The pool should also be used for any
    /// temporary buffers created during execution.
    fn run_in_place(
        &self,
        #[allow(unused)] in_place: InPlaceInputs,
        #[allow(unused)] ctx: &OpRunContext,
    ) -> Result<OutputList, OpError> {
        Err(OpError::InvalidValue("In-place execution not supported"))
    }

    /// Return the IDs of inputs which can be pre-packed using [`prepack`](Operator::prepack).
    fn prepack_inputs(&self) -> SmallVec<[usize; 1]> {
        SmallVec::new()
    }

    /// Pre-pack an input for more efficient inference later.
    ///
    /// `index` specifies the input ID and should be one of the inputs returned
    /// by [`prepack_inputs`](Operator::prepack_inputs).
    fn prepack(
        &self,
        #[allow(unused)] index: usize,
        #[allow(unused)] input: ValueView,
    ) -> Option<PrepackedInput> {
        None
    }

    /// Return the [`SubgraphOperator`] implementation for this operator, if
    /// this operator has subgraphs.
    fn as_subgraph_op(&self) -> Option<&dyn SubgraphOperator> {
        None
    }

    /// Return the shape inference implementation for this operator.
    ///
    /// The implementation may return `None` if shape inference is unsupported.
    fn as_infer_shapes(&self) -> Option<&dyn InferShapes>;
}

impl dyn Operator {
    /// Downcast this operator to a concrete type.
    pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
        (self as &dyn Any).downcast_ref()
    }
}

impl dyn Operator + Send + Sync {
    /// Downcast this operator to a concrete type.
    pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
        (self as &dyn Any).downcast_ref()
    }
}

/// Trait for operators which contain subgraphs, such as `If`, `Loop` etc.
pub trait SubgraphOperator: Operator {
    /// Return a list of subgraphs used by this operator.
    fn subgraphs(&self) -> SmallVec<[&Graph; 2]>;

    /// Execute the operator with the given inputs and captured values.
    ///
    /// This should be used instead of [`Operator::run`] for operators that
    /// implement this trait.
    fn run_subgraph<'a>(
        &'a self,
        ctx: &OpRunContext,
        #[allow(unused)] captures: CaptureEnv,
        #[allow(unused)] weight_cache: Option<&[WeightCache]>,
        #[allow(unused)] profiler: Option<&mut Profiler<'a>>,
        #[allow(unused)] run_opts: Option<RunOptions>,
    ) -> Result<OutputList, RunError>;
}

/// Convenience methods that make it easier to run operators in tests.
#[cfg(test)]
pub trait OperatorExt: Operator {
    /// Run an operator and extract the first output as a tensor with a given
    /// type.
    ///
    /// `inputs` is a tuple of tensor references or other values that can be
    /// converted to [`ValueView`].
    fn run_simple<'a, I: Into<InputList<'a>>, O: TryFrom<Value>>(
        &self,
        inputs: I,
    ) -> Result<O, OpError>
    where
        OpError: From<<O as TryFrom<Value>>::Error>,
    {
        let pool = BufferPool::new();
        let inputs = inputs.into();
        let ctx = OpRunContext::new(&pool, &inputs, OutputMask::all_used(1));
        let mut outputs = self.run(&ctx)?;
        Ok(outputs.remove(0).try_into()?)
    }

    /// Run an operator with a mutable input and extract the first output.
    fn run_simple_in_place<'a, M: Into<Value>, I: Into<InputList<'a>>, O: TryFrom<Value>>(
        &self,
        mut_input: M,
        inputs: I,
    ) -> Result<O, OpError>
    where
        OpError: From<<O as TryFrom<Value>>::Error>,
    {
        let pool = BufferPool::new();
        let inputs = inputs.into();
        let ctx = OpRunContext::new(&pool, &inputs, OutputMask::all_used(1));
        let in_place = InPlaceInputs::from((0, mut_input.into()));
        let mut outputs = self.run_in_place(in_place, &ctx)?;
        let typed_output = outputs.remove(0).try_into()?;
        Ok(typed_output)
    }
}

#[cfg(test)]
impl<O: ?Sized + Operator> OperatorExt for O {}

/// List of inputs for an operator evaluation.
///
/// This is an owned or borrowed collection of `Option<ValueView>`s with methods
/// to conveniently extract inputs and produce appropriate errors if inputs are
/// missing or of the wrong type.
///
/// An InputList can be constructed from tuples of `impl Into<ValueView>` types
/// (eg. `TensorView`, `&Tensor`) via `Into`. It can also be created or
/// extended from iterators of `ValueView`s or `Option<ValueView>`s.
#[derive(Clone)]
pub struct InputList<'a> {
    inputs: Cow<'a, [Option<ValueView<'a>>]>,

    /// Callback that retrieves the pre-packed copy of an input with a given
    /// index.
    get_prepacked: Option<&'a dyn Fn(usize) -> Option<&'a PrepackedInput>>,
}

impl<'a> InputList<'a> {
    /// Construct an empty input list.
    pub fn new() -> InputList<'a> {
        InputList {
            inputs: Cow::Owned(vec![]),
            get_prepacked: None,
        }
    }

    pub fn len(&self) -> usize {
        self.inputs.len()
    }

    pub fn is_empty(&self) -> bool {
        self.inputs.is_empty()
    }

    /// Append an input to the list.
    ///
    /// This will copy the existing inputs into a new owned vector.
    pub fn push<I: Into<ValueView<'a>>>(&mut self, inp: I) {
        self.inputs.to_mut().push(Some(inp.into()))
    }

    /// Append an optional input to the list.
    ///
    /// This will copy the existing inputs into a new owned vector.
    pub fn push_optional<I: Into<ValueView<'a>>>(&mut self, inp: Option<I>) {
        self.inputs.to_mut().push(inp.map(|inp| inp.into()))
    }

    /// Construct an input list from a slice of non-optional inputs.
    ///
    /// This copies the inputs into a new vector of `Option<ValueView>`s. Using
    /// [`from_optional`](Self::from_optional) is more efficient.
    pub fn from(inputs: &[ValueView<'a>]) -> InputList<'a> {
        InputList {
            inputs: inputs.iter().cloned().map(Some).collect(),
            get_prepacked: None,
        }
    }

    /// Construct an input list from a slice of optional inputs.
    ///
    /// This is a cheap conversion that borrows `inputs`.
    pub fn from_optional(inputs: &'a [Option<ValueView<'a>>]) -> InputList<'a> {
        InputList {
            inputs: Cow::Borrowed(inputs),
            get_prepacked: None,
        }
    }

    /// Configure a callback that will get or create a pre-packed copy of the
    /// input with a given index.
    pub fn with_prepacked(
        mut self,
        lookup: &'a dyn Fn(usize) -> Option<&'a PrepackedInput>,
    ) -> Self {
        self.get_prepacked = Some(lookup);
        self
    }

    /// Get an optional input.
    pub fn get(&self, index: usize) -> Option<ValueView<'a>> {
        self.inputs.get(index).cloned().flatten()
    }

    /// Get the pre-packed version of a weight input, if available.
    pub fn get_prepacked(&self, index: usize) -> Option<&'a PrepackedInput> {
        self.get_prepacked.and_then(|gp| gp(index))
    }

    /// Get a mutable reference to an input.
    ///
    /// This will convert the list into an owned list of inputs first.
    pub fn get_mut(&mut self, index: usize) -> Option<&mut ValueView<'a>> {
        self.inputs.to_mut().get_mut(index)?.as_mut()
    }

    /// Convert an optional input into a tensor or scalar.
    pub fn get_as<T>(&self, index: usize) -> Result<Option<T>, OpError>
    where
        T: TryFrom<ValueView<'a>, Error = TryFromValueError>,
    {
        self.get(index)
            .map(|input| {
                input
                    .try_into()
                    .map_err(|error| OpError::InputCastFailed { index, error })
            })
            .transpose()
    }

    /// Get a required operator input.
    pub fn require(&self, index: usize) -> Result<ValueView<'a>, OpError> {
        self.get(index).ok_or(OpError::MissingInputs)
    }

    /// Convert a required input into a tensor or scalar.
    pub fn require_as<T>(&self, index: usize) -> Result<T, OpError>
    where
        T: TryFrom<ValueView<'a>, Error = TryFromValueError>,
    {
        self.require(index).and_then(|input| {
            input
                .try_into()
                .map_err(|error| OpError::InputCastFailed { index, error })
        })
    }

    /// Return the index of the first input which is present (not `None`).
    ///
    /// This is useful for commutative operators running in-place, where the
    /// single non-in-place input may be at a different position depending on
    /// which input was selected for in-place execution.
    pub fn first_present(&self) -> Option<usize> {
        self.inputs.iter().position(|input| input.is_some())
    }

    /// Convert the [`first_present`](Self::first_present) input into a tensor
    /// or scalar.
    pub fn require_first_present_as<T>(&self) -> Result<T, OpError>
    where
        T: TryFrom<ValueView<'a>, Error = TryFromValueError>,
    {
        let index = self.first_present().ok_or(OpError::MissingInputs)?;
        self.require_as(index)
    }

    /// Return an iterator over provided inputs.
    ///
    /// Use [`Iterator::flatten`] to skip missing optional inputs.
    pub fn iter<'b>(&'b self) -> impl Iterator<Item = Option<ValueView<'a>>> + 'b {
        self.inputs.iter().cloned()
    }
}

impl Default for InputList<'_> {
    fn default() -> Self {
        Self::new()
    }
}

impl<'a, I: Into<ValueView<'a>>> From<I> for InputList<'a> {
    fn from(val: I) -> InputList<'a> {
        InputList::from(&[val.into()])
    }
}

impl<'a> From<()> for InputList<'a> {
    fn from(_: ()) -> InputList<'a> {
        Self::default()
    }
}

impl<'a, I1: Into<ValueView<'a>>> From<(I1,)> for InputList<'a> {
    fn from((a,): (I1,)) -> InputList<'a> {
        InputList::from(&[a.into()])
    }
}

impl<'a, I1: Into<ValueView<'a>>, I2: Into<ValueView<'a>>> From<(I1, I2)> for InputList<'a> {
    fn from((a, b): (I1, I2)) -> InputList<'a> {
        InputList::from(&[a.into(), b.into()])
    }
}

impl<'a, I1: Into<ValueView<'a>>, I2: Into<ValueView<'a>>, I3: Into<ValueView<'a>>>
    From<(I1, I2, I3)> for InputList<'a>
{
    fn from((a, b, c): (I1, I2, I3)) -> InputList<'a> {
        InputList::from(&[a.into(), b.into(), c.into()])
    }
}

impl<
    'a,
    I1: Into<ValueView<'a>>,
    I2: Into<ValueView<'a>>,
    I3: Into<ValueView<'a>>,
    I4: Into<ValueView<'a>>,
> From<(I1, I2, I3, I4)> for InputList<'a>
{
    fn from((a, b, c, d): (I1, I2, I3, I4)) -> InputList<'a> {
        InputList::from(&[a.into(), b.into(), c.into(), d.into()])
    }
}

impl<
    'a,
    I1: Into<ValueView<'a>>,
    I2: Into<ValueView<'a>>,
    I3: Into<ValueView<'a>>,
    I4: Into<ValueView<'a>>,
    I5: Into<ValueView<'a>>,
> From<(I1, I2, I3, I4, I5)> for InputList<'a>
{
    fn from((a, b, c, d, e): (I1, I2, I3, I4, I5)) -> InputList<'a> {
        InputList::from(&[a.into(), b.into(), c.into(), d.into(), e.into()])
    }
}

impl<'a> Extend<ValueView<'a>> for InputList<'a> {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = ValueView<'a>>,
    {
        for item in iter {
            self.push(item);
        }
    }
}

impl<'a> Extend<Option<ValueView<'a>>> for InputList<'a> {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = Option<ValueView<'a>>>,
    {
        for item in iter {
            self.push_optional(item);
        }
    }
}

impl<'a, A> FromIterator<A> for InputList<'a>
where
    InputList<'a>: Extend<A>,
{
    fn from_iter<T>(iter: T) -> Self
    where
        T: IntoIterator<Item = A>,
    {
        let mut list = InputList::new();
        list.extend(iter);
        list
    }
}

#[cfg(test)]
mod tests {
    use rten_base::bit_set::BitSet;
    use rten_tensor::prelude::*;
    use rten_tensor::{Tensor, TensorView};

    use crate::operator::{InputList, OpError, Operator, OutputMask};
    use crate::ops::{Add, Sub};

    #[test]
    fn test_output_mask() {
        // Lengths below, at and above the number of tracked outputs.
        for len in [5, 32, 33, 100] {
            let mask = OutputMask::all_used(len);
            assert_eq!(mask.len(), len);
            for i in 0..len + 10 {
                assert_eq!(mask.is_used(i), i < len);
            }
        }

        // Positions beyond the first 32 are always treated as used, even if
        // not marked as used in the mask.
        let mask = OutputMask::new(BitSet::from_indices([0, 3]), 100);
        for i in 0..100 {
            assert_eq!(mask.is_used(i), i == 0 || i == 3 || i >= 32);
        }
    }

    #[test]
    fn test_input_list_require_as_error_index() {
        let tensor = Tensor::<f32>::zeros(&[2, 2]);

        let inputs = InputList::from(&[tensor.view().into(), tensor.view().into()]);
        let err = inputs.require_as::<TensorView<i32>>(1).err().unwrap();
        assert!(matches!(err, OpError::InputCastFailed { index: 1, .. }));
    }

    #[test]
    fn test_downcast_operator() {
        let add_op = Add {};
        let sub_op = Sub {};

        let add_op_dyn: &dyn Operator = &add_op;
        let sub_op_dyn: &dyn Operator = &sub_op;

        assert!(add_op_dyn.downcast_ref::<Add>().is_some());
        assert!(sub_op_dyn.downcast_ref::<Sub>().is_some());
    }
}