tokitai-operator 0.1.0

//! Shape-manipulation operators.
//!
//! Each op takes one or more dense tensors and produces a dense tensor
//! with a different shape, same underlying element order (Reshape /
//! Permute / Transpose / Slice) or different element order (Concat /
//! Broadcast). The shape ops never change the element type or the
//! mathematical domain; they only change layout metadata. As a result,
//! they are backend-neutral and lowering is purely a data-movement
//! operation that preserves exact i64 values.
//!
//! Permute vs. Transpose:
//!   - `Transpose` swaps exactly two axes; rank is preserved.
//!   - `Permute` takes a full permutation vector of length `rank`.
//!     `Transpose` is a strict subset of `Permute`; we keep both as
//!     separate ops so the planner can pick the cheaper one when only
//!     two axes are involved.
//!
//! Slice semantics:
//!   - `Slice` uses the Python half-open convention `start..end`
//!     (the slice covers elements at indices `start, start+1, ...,
//!     end-1`). An end of `usize::MAX` means "to the end of the
//!     axis".
//!
//! Concat semantics:
//!   - `Concat` accepts 2 or more input tensors. All inputs must
//!     have the same rank, and the same shape on every axis except
//!     the concatenation axis. Output rank equals input rank.

use crate::domain::{Claim, Contract, ContractId, ContractSet, Evidence, Scope};
use crate::object::{Dim, ObjectKind, ObjectMeta, Shape};
use crate::{Error, Result};

use super::{LayerBehavior, OpInput, OpOutput, OpSignature, Operator};

// ---------------------------------------------------------------------------
// Small helper: deterministic contracts for shape-preserving data movement.
// ---------------------------------------------------------------------------
fn shape_arithmetic_contracts(scope: Scope) -> ContractSet {
    ContractSet::from_iter([
        Contract::new(
            ContractId(30),
            Claim::Deterministic,
            scope.clone(),
            Evidence::Axiom,
        ),
        Contract::new(ContractId(31), Claim::Exact, scope, Evidence::Axiom),
    ])
}

fn tensor_input_check(op: &str, inputs: &[ObjectMeta], expected: usize) -> Result<()> {
    if inputs.len() != expected {
        return Err(Error::operator(format!(
            "{op} expects {expected} input(s), got {}",
            inputs.len()
        )));
    }
    for (i, m) in inputs.iter().enumerate() {
        if m.object_kind != ObjectKind::Tensor {
            return Err(Error::operator(format!(
                "{op} only supports tensor inputs, input {i} is {:?}",
                m.object_kind
            )));
        }
    }
    Ok(())
}

fn parse_target_shape(meta: &ObjectMeta) -> Result<Vec<Dim>> {
    if meta.object_kind != ObjectKind::Tensor {
        return Err(Error::operator(format!(
            "target shape input must be a tensor, got {:?}",
            meta.object_kind
        )));
    }
    Ok(meta.shape.dims.clone())
}

// ---------------------------------------------------------------------------
// Reshape: out = input reshaped to a target shape (element count must match).
// ---------------------------------------------------------------------------

/// Reshape a tensor to a new shape. Element count must match exactly
/// (when both source and target shapes are fully static).
#[derive(Debug, Clone, Copy, Default)]
pub struct ReshapeOp;

impl Operator for ReshapeOp {
    fn name(&self) -> &'static str {
        "reshape"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "target_shape".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let target_dims = parse_target_shape(&inputs[1])?;
        let mut output = inputs[0].clone();
        output.shape = Shape::new(target_dims);
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Transpose: swap exactly two axes.
// ---------------------------------------------------------------------------

/// Transpose (swap two axes). The two axes are read from a 2-element
/// i32 tensor input (axis0, axis1). The lowerings support negative
/// indexing.
#[derive(Debug, Clone, Copy, Default)]
pub struct TransposeOp;

impl Operator for TransposeOp {
    fn name(&self) -> &'static str {
        "transpose"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "axes".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let rank = inputs[0].shape.rank();
        let mut output = inputs[0].clone();
        // We don't know the axis values at type-erased plan time, so
        // we keep the same rank but record the op as a generic
        // permutation. The lowering will re-order the data
        // according to the runtime axes tensor.
        output.shape = Shape::new(
            (0..rank)
                .map(|i| inputs[0].shape.dims[i].clone())
                .collect::<Vec<Dim>>(),
        );
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Permute: arbitrary permutation of axes (rank is preserved).
// ---------------------------------------------------------------------------

/// Permute: take a permutation vector of length `rank` and reorder
/// the axes accordingly. The permutation is read from a 1-D i32
/// tensor input.
#[derive(Debug, Clone, Copy, Default)]
pub struct PermuteOp;

impl Operator for PermuteOp {
    fn name(&self) -> &'static str {
        "permute"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "permutation".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let rank = inputs[0].shape.rank();
        let mut output = inputs[0].clone();
        output.shape = Shape::new(
            (0..rank)
                .map(|i| inputs[0].shape.dims[i].clone())
                .collect::<Vec<Dim>>(),
        );
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Slice: out = input[start..end] along `axis` (Python half-open).
// ---------------------------------------------------------------------------

/// Slice along a single axis using the Python half-open convention
/// `start..end` (indices `start, start+1, ..., end-1`). The `end`
/// value `usize::MAX` means "to the end of the axis". The `axis`
/// and (start, end) bounds are read from a 3-element i32 tensor
/// input: `[axis, start, end]`.
#[derive(Debug, Clone, Copy, Default)]
pub struct SliceOp;

impl Operator for SliceOp {
    fn name(&self) -> &'static str {
        "slice"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "bounds".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let mut output = inputs[0].clone();
        // The actual rank/dim slicing happens at lowering time once
        // we have the runtime bounds. We keep the same rank; the
        // lowering shrinks the affected dim.
        let rank = inputs[0].shape.rank();
        output.shape = Shape::new(
            (0..rank)
                .map(|i| inputs[0].shape.dims[i].clone())
                .collect::<Vec<Dim>>(),
        );
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Concat: stack N tensors along one axis.
// ---------------------------------------------------------------------------

/// Concat 2 or more tensors along a single axis. All input tensors
/// must have the same rank and the same shape on every axis except
/// the concatenation axis.
#[derive(Debug, Clone, Copy, Default)]
pub struct ConcatOp;

impl Operator for ConcatOp {
    fn name(&self) -> &'static str {
        "concat"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![OpInput {
                name: "first".to_string(),
            }],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        if inputs.is_empty() {
            return Err(Error::operator("concat expects at least 1 input"));
        }
        for (i, m) in inputs.iter().enumerate() {
            if m.object_kind != ObjectKind::Tensor {
                return Err(Error::operator(format!(
                    "concat only supports tensor inputs, input {i} is {:?}",
                    m.object_kind
                )));
            }
        }
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Broadcast: explicitly broadcast an input to a target shape.
// ---------------------------------------------------------------------------

/// Broadcast: explicitly expand an input tensor to a target shape.
/// This is the explicit (op-form) counterpart of the implicit
/// broadcasting that binary arithmetic ops do. The target shape is
/// read from a second tensor input.
#[derive(Debug, Clone, Copy, Default)]
pub struct BroadcastOp;

impl Operator for BroadcastOp {
    fn name(&self) -> &'static str {
        "broadcast"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "target_shape".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let target_dims = parse_target_shape(&inputs[1])?;
        let mut output = inputs[0].clone();
        output.shape = Shape::new(target_dims);
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Flatten: collapse every dimension into a single 1-D dim.
// ---------------------------------------------------------------------------

/// Flatten: collapse every dimension of an N-D tensor into a single
/// 1-D dimension. A rank-0 (scalar) tensor is mapped to `[1]` (this
/// matches the `Tensor::try_from_vec` convention where rank-0 has
/// product 1).
#[derive(Debug, Clone, Copy, Default)]
pub struct FlattenOp;

impl Operator for FlattenOp {
    fn name(&self) -> &'static str {
        "flatten"
    }

    fn signature(&self) -> OpSignature {
        unary_signature_like()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        if inputs.len() != 1 {
            return Err(Error::operator(format!(
                "flatten expects 1 input, got {}",
                inputs.len()
            )));
        }
        let m = &inputs[0];
        if m.object_kind != ObjectKind::Tensor {
            return Err(Error::operator(format!(
                "flatten only supports tensor inputs, got {:?}",
                m.object_kind
            )));
        }
        let mut output = m.clone();
        // The actual flattened dim size is determined at lowering
        // time. We use a DataDependent placeholder when the input
        // rank is 0 (the canonical "scalar -> [1]" rule) or
        // Symbolic/Bounded dims that prevent a static product.
        let product: Option<usize> = if m.shape.dims.is_empty() {
            Some(1)
        } else {
            m.shape
                .dims
                .iter()
                .try_fold(1usize, |acc, d| d.value().map(|v| acc * v))
        };
        let flat_dim = match product {
            Some(n) => Dim::Static(n),
            None => Dim::DataDependent("flatten_product".to_string()),
        };
        output.shape = Shape::new(vec![flat_dim]);
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Squeeze: drop all size-1 dims (or a specific subset, declared as a
// runtime mask tensor in input[1]).
// ---------------------------------------------------------------------------

/// Squeeze: drop size-1 dimensions. When no `dims` input is given
/// (only the data tensor), all size-1 dims are dropped. When a
/// `dims` mask tensor is given, only the dims whose index appears
/// in the mask (and which are size-1) are dropped.
#[derive(Debug, Clone, Copy, Default)]
pub struct SqueezeOp;

impl Operator for SqueezeOp {
    fn name(&self) -> &'static str {
        "squeeze"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![OpInput {
                name: "input".to_string(),
            }],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        if inputs.len() != 1 {
            return Err(Error::operator(format!(
                "squeeze expects 1 input, got {}",
                inputs.len()
            )));
        }
        let m = &inputs[0];
        if m.object_kind != ObjectKind::Tensor {
            return Err(Error::operator(format!(
                "squeeze only supports tensor inputs, got {:?}",
                m.object_kind
            )));
        }
        let mut output = m.clone();
        // Eagerly drop all size-1 dims we can see at plan time. The
        // runtime-conditional cases (Symbolic/Bounded dims) keep
        // their dim and the lowering handles them.
        output.shape = Shape::new(
            m.shape
                .dims
                .iter()
                .filter(|d| !matches!(d, Dim::Static(1)))
                .cloned()
                .collect::<Vec<Dim>>(),
        );
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// Unsqueeze: insert a size-1 dim at the given axis (i32 from a 1-elem
// tensor).
// ---------------------------------------------------------------------------

/// Unsqueeze: insert a single size-1 dimension at the given axis.
/// The axis is read from a 1-element i32 tensor input. Negative axes
/// count from the end.
#[derive(Debug, Clone, Copy, Default)]
pub struct UnsqueezeOp;

impl Operator for UnsqueezeOp {
    fn name(&self) -> &'static str {
        "unsqueeze"
    }

    fn signature(&self) -> OpSignature {
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "axis".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 2)?;
        let rank = inputs[0].shape.rank();
        let mut output = inputs[0].clone();
        let mut new_dims = inputs[0].shape.dims.clone();
        new_dims.insert(rank, Dim::Static(1));
        output.shape = Shape::new(new_dims);
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        shape_arithmetic_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

fn unary_signature_like() -> OpSignature {
    OpSignature {
        inputs: vec![OpInput {
            name: "input".to_string(),
        }],
        outputs: vec![OpOutput {
            name: "out".to_string(),
        }],
    }
}