topos 0.13.0

use std::fmt;
use std::ops::{Add, Div, Mul, Neg, Sub};
use std::sync::Arc;

use static_assertions::assert_impl_all;

use crate::backend::MapTask;

use super::elementary::MapOperation;
use super::gemm;
use super::layout::{Layout, Strides};
use super::normalized::BatchNormTask;
use super::recordable::{composed_batch_norm, composed_max_pool, composed_windowed_patches};
use super::storage::Storage;
use super::{Differentiable, Element, Elementary, Recordable, Shape};

// Entry-time thread-safety contract; the anchor rationale is documented
// in `network.rs`.
assert_impl_all!(Tensor<f64>: Send, Sync);

/// A dense tensor with an immutable, runtime-defined [`Shape`] and a shared
/// element buffer read through a strided layout.
///
/// The elements are held behind a `Storage` representation: an
/// `Arc`-shared row-major buffer addressed by strides and an offset, or a
/// non-allocating constant. Cloning shares the buffer and clones only the
/// metadata; it does not clone the elements. Because tensors are immutable
/// and buffer-shared, view operations that alias a buffer (transpose and
/// broadcast) are always safe: no operation ever writes through an alias.
///
/// Arithmetic and the elementwise maps operate in logical row-major
/// order. Binary elementwise operations require identical shapes and
/// never broadcast implicitly. Broadcasting is available only through
/// [`broadcast_like`](Tensor::broadcast_like) and
/// [`broadcast_along`](Tensor::broadcast_along), and it produces a
/// view rather than copying.
///
/// [`matmul`](Tensor::matmul) requires rank-2 or batched higher-rank
/// operands, and [`transpose`](Tensor::transpose) accepts ranks 0
/// through 2, returning a view. Reductions, explicit broadcasts, and
/// reshaping are rank-general.
#[derive(Debug, Clone)]
pub struct Tensor<Element> {
    storage: Storage<Element>,
}

impl<Element> Tensor<Element> {
    /// Returns the shape of this tensor: its extent along every axis.
    ///
    /// It is what record-time shape inference seeds leaves with. A
    /// scalar is rank 0.
    pub fn shape(&self) -> Shape {
        self.logical_shape().clone()
    }

    /// Returns the logical shape, the one descriptor every representation
    /// answers for.
    fn logical_shape(&self) -> &Shape {
        match &self.storage {
            Storage::Dense { layout, .. } => layout.shape(),
            Storage::Constant { shape, .. } => shape,
            Storage::Selection { shape, .. } => shape,
        }
    }

    /// Returns the repeated value when this tensor is stored as a
    /// non-allocating constant fill.
    ///
    /// This is a representation fact, not a mathematical predicate:
    /// a dense buffer that happens to hold one value throughout
    /// answers `None`. Displays and norms use it as an O(1) fast
    /// path over fills such as gradient seeds and counted leaves.
    pub fn as_constant(&self) -> Option<&Element> {
        match &self.storage {
            Storage::Constant { value, .. } => Some(value),
            _ => None,
        }
    }

    /// Returns the row indices of a `Selection` payload.
    ///
    /// # Panics
    /// Panics if `self` is not a selection built with [`Tensor::selection`].
    fn selection_indices(&self) -> &[usize] {
        match &self.storage {
            Storage::Selection { indices, .. } => indices,
            _ => panic!("gather requires a selection tensor built with `Tensor::selection`"),
        }
    }

    /// Returns this payload as a gemm operand when it is a dense
    /// tensor of rank two or higher: the backing slice from the
    /// layout's offset, with the layout's full strides — the trailing
    /// two drive each `GemmTask`, the leading ones step the batch.
    /// Other storages answer `None` and take the logical path.
    fn gemm_operand(&self) -> Option<(&[Element], &[usize])> {
        match &self.storage {
            Storage::Dense { data, layout } if layout.rank() >= 2 => {
                Some((&data.as_slice()[layout.offset()..], layout.strides()))
            }
            _ => None,
        }
    }

    /// Returns the elements as a contiguous slice when the tensor is stored
    /// as a contiguous dense buffer, or `None` for a strided view or a
    /// constant.
    pub fn as_slice(&self) -> Option<&[Element]> {
        match &self.storage {
            Storage::Dense { data, layout } if layout.is_contiguous() => {
                let start = layout.offset();
                Some(&data.as_slice()[start..start + layout.volume()])
            }
            _ => None,
        }
    }

    /// Returns the buffer window a strided dense view addresses together
    /// with its layout rebased to the window start, or `None` when the
    /// storage is not a strided view or the window is not strictly narrower
    /// than the volume: only a broadcast view computes fewer elements and
    /// earns staying a view. An equal-width window (a transpose) must keep
    /// the logical scalar walk — the backend seam reading the window would
    /// silently replace the documented bitwise fallback — and a wider one
    /// (a narrow sliver) would cost more than the walk it replaces.
    ///
    /// An elementwise transform commutes with any view, so a caller may
    /// transform the window and keep the layout instead of materializing
    /// the logical order. The window can include buffer elements the view
    /// skips; the transform visits them, and no logical access reads them.
    fn strided_window(&self) -> Option<(&[Element], Layout)> {
        let Storage::Dense { data, layout } = &self.storage else {
            return None;
        };
        if layout.is_contiguous() || layout.span() >= layout.volume() {
            return None;
        }
        let start = layout.offset();
        Some((
            &data.as_slice()[start..start + layout.span()],
            layout.rebased(),
        ))
    }
}

impl<Element: Clone> Tensor<Element> {
    /// Returns the elements in logical row-major order.
    ///
    /// A contiguous dense buffer iterates its slice directly; a strided
    /// view walks its layout with an odometer; a constant repeats its
    /// single value across the shape's volume.
    ///
    /// It yields owned elements rather than references because a
    /// representation that computes its elements has nothing to lend,
    /// and for the numeric payloads a clone is the same load a
    /// borrow-then-copy compiles to.
    pub fn iter(&self) -> impl Iterator<Item = Element> + '_ {
        match &self.storage {
            Storage::Constant { shape, value } => ElementIter::Constant {
                value,
                remaining: shape.volume(),
            },
            Storage::Dense { data, layout } if layout.is_contiguous() => {
                let start = layout.offset();
                ElementIter::Contiguous(data.as_slice()[start..start + layout.volume()].iter())
            }
            Storage::Dense { data, layout } => ElementIter::Strided {
                data: data.as_slice(),
                shape: layout.shape().axes(),
                strides: layout.strides(),
                coordinates: std::iter::repeat_n(0usize, layout.rank()).collect(),
                index: layout.offset(),
                remaining: layout.volume(),
            },
            Storage::Selection {
                indices,
                shape,
                zero,
                one,
            } => ElementIter::Selection {
                indices: indices.as_slice(),
                vocab: shape.axes()[1],
                zero,
                one,
                position: 0,
                total: shape.volume(),
            },
        }
    }

    /// Returns the elements in logical row-major order as an owned vector.
    pub fn to_vec(&self) -> Vec<Element> {
        self.iter().collect()
    }

    /// Returns this tensor with every element converted into `Target`
    /// through its `From` conversion, preserving the storage
    /// representation: a constant stays a constant, a selection stays
    /// a selection, and a dense view keeps its layout, so a broadcast
    /// converts only its distinct buffer elements.
    ///
    /// It is the precision boundary for mixed-precision work —
    /// loading an `f32` checkpoint into a `Tensor<Bf16>` model, or
    /// widening bf16 results back — priced at one conversion per
    /// stored element.
    pub fn convert<Target: From<Element>>(&self) -> Tensor<Target> {
        match &self.storage {
            Storage::Dense { data, layout } => Tensor {
                storage: Storage::Dense {
                    data: Arc::new(data.iter().cloned().map(Target::from).collect()),
                    layout: layout.clone(),
                },
            },
            Storage::Constant { shape, value } => Tensor {
                storage: Storage::Constant {
                    shape: shape.clone(),
                    value: Target::from(value.clone()),
                },
            },
            Storage::Selection {
                indices,
                shape,
                zero,
                one,
            } => Tensor {
                storage: Storage::Selection {
                    indices: Arc::clone(indices),
                    shape: shape.clone(),
                    zero: Target::from(zero.clone()),
                    one: Target::from(one.clone()),
                },
            },
        }
    }

    /// Returns the rank-0 tensor's single element: the scalar
    /// projection, and the read-back the teaching examples end on.
    ///
    /// It is deliberately loud — one call, rank-checked — rather than
    /// a `Deref` to the element: a silent projection is the kind of
    /// magic the explicitness rule forbids. A rank-1 tensor of one
    /// element does not qualify; reshape it first.
    ///
    /// # Panics
    /// Panics if this tensor is not rank 0.
    pub fn scalar(&self) -> Element {
        assert_eq!(
            self.logical_shape().rank(),
            0,
            "scalar reads a rank-0 tensor, got {}",
            self.logical_shape()
        );
        self.get(0)
    }

    /// Returns the element at logical row-major `position`.
    ///
    /// It is the general per-element read shared by every operation; a
    /// dense layout resolves it through the stride and offset arithmetic,
    /// and a constant answers with its single value.
    ///
    /// It returns an owned element rather than a reference because a
    /// representation that computes its elements has nothing to lend,
    /// and because nearly every caller wants the value: for the numeric
    /// payloads a clone is the same load a borrow-then-copy compiles to.
    fn get(&self, position: usize) -> Element {
        match &self.storage {
            Storage::Dense { data, layout } => data[layout.storage_index(position)].clone(),
            Storage::Constant { value, .. } => value.clone(),
            Storage::Selection {
                indices,
                shape,
                zero,
                one,
            } => {
                let vocab = shape.axes()[1];
                if indices[position / vocab] == position % vocab {
                    one.clone()
                } else {
                    zero.clone()
                }
            }
        }
    }
}

impl<Element: Differentiable> Tensor<Element> {
    /// Creates a tensor of `shape` from `elements` in row-major order.
    ///
    /// # Panics
    /// Panics if the shape's volume overflows `usize`, the number of elements
    /// differs from that volume, or the shape holds no elements. Empty tensors
    /// are unsupported because reductions initialize their accumulator from
    /// an existing element.
    pub fn new(shape: impl Into<Shape>, elements: impl Into<Vec<Element>>) -> Self {
        Self::dense(shape.into(), elements.into())
    }

    /// Creates a tensor of `shape` with every element set to `element`,
    /// stored as a non-allocating constant.
    ///
    /// # Panics
    /// Panics if the shape's volume overflows `usize` or the shape holds no
    /// elements, as documented on [`Tensor::new`].
    pub fn filled(shape: impl Into<Shape>, element: Element) -> Self {
        Self::constant(shape.into(), element)
    }

    /// Builds a contiguous dense tensor of `shape` from `elements`.
    ///
    /// Every dense tensor is created here, so this is where the tensor
    /// invariant is proven: the shape's volume is representable (checked
    /// by [`Shape::volume`]), positive, and equal to the buffer length.
    fn dense(shape: Shape, elements: Vec<Element>) -> Self {
        assert_eq!(
            shape.volume(),
            elements.len(),
            "tensor shape does not match its number of elements"
        );
        assert!(
            !elements.is_empty(),
            "tensors must hold at least one element"
        );
        Self {
            storage: Storage::Dense {
                layout: Layout::contiguous(shape),
                data: Arc::new(elements),
            },
        }
    }

    /// Builds a constant tensor of `shape` filled with `value`.
    ///
    /// Every constant tensor is created here, so the tensor invariant is
    /// proven here as on [`Tensor::dense`]: a representable, positive
    /// volume.
    fn constant(shape: Shape, value: Element) -> Self {
        assert!(shape.volume() > 0, "tensors must hold at least one element");
        Self {
            storage: Storage::Constant { shape, value },
        }
    }

    /// Creates the one-hot `[indices.len(), vocab]` selection matrix whose
    /// row `i` is `one` at column `indices[i]` and zero elsewhere, stored as
    /// its indices rather than a dense buffer.
    ///
    /// It carries the token indices of an embedding lookup: feed it as a
    /// per-run input and read it with [`Recordable::gather`](super::Recordable::gather).
    /// `one` is the value placed at each selected position (the
    /// multiplicative identity, e.g. `1.0`); the zero is derived from it.
    ///
    /// # Panics
    /// Panics if `vocab` is zero, `indices` is empty, any index is not
    /// below `vocab`, or the `[indices.len(), vocab]` volume overflows
    /// `usize`.
    pub fn selection(indices: impl Into<Vec<usize>>, vocab: usize, one: Element) -> Self {
        let indices = indices.into();
        assert!(vocab > 0, "a selection needs a non-empty vocabulary");
        assert!(
            !indices.is_empty(),
            "tensors must hold at least one element"
        );
        assert!(
            indices.len().checked_mul(vocab).is_some(),
            "shape volume overflows `usize`"
        );
        for &index in &indices {
            assert!(
                index < vocab,
                "selection index {index} is out of vocabulary {vocab}"
            );
        }
        let zero = Element::zero();
        let shape = Shape::new([indices.len(), vocab]);
        Self {
            storage: Storage::Selection {
                indices: Arc::new(indices),
                shape,
                zero,
                one,
            },
        }
    }

    /// Returns an equivalent contiguous dense tensor, materializing any
    /// non-dense or strided representation.
    ///
    /// It is the correctness fallback for view operations that a `Selection`
    /// does not model directly (transpose, permute, narrow, axis broadcast):
    /// densify first, then take the dense view.
    fn densify(&self) -> Self {
        match &self.storage {
            Storage::Dense { layout, .. } if layout.is_contiguous() => self.clone(),
            _ => Self::dense(self.logical_shape().clone(), self.to_vec()),
        }
    }

    /// Returns a tensor with every element passed through `transform`.
    ///
    /// A constant maps in place to another constant; a contiguous
    /// dense buffer maps over its slice directly — bypassing the
    /// per-element iterator dispatch, which measured well below
    /// memory speed; a strided view maps its buffer window and keeps
    /// the layout, so a broadcast view transforms only its distinct
    /// elements and stays a view; and everything else materializes
    /// through the logical-order iterator. A pure `transform` gives
    /// every logical element the same value on every path, so the
    /// lanes agree bitwise.
    fn map(&self, transform: impl Fn(&Element) -> Element) -> Self {
        if let Storage::Constant { shape, value } = &self.storage {
            return Self::constant(shape.clone(), transform(value));
        }
        if let Some(elements) = self.as_slice() {
            return Self::dense(
                self.logical_shape().clone(),
                elements.iter().map(transform).collect(),
            );
        }
        if let Some((window, layout)) = self.strided_window() {
            return Self {
                storage: Storage::Dense {
                    data: Arc::new(window.iter().map(transform).collect()),
                    layout,
                },
            };
        }
        Self::dense(
            self.logical_shape().clone(),
            self.iter().map(|element| transform(&element)).collect(),
        )
    }

    /// Combines two tensors element by element with `combine`.
    ///
    /// Two constants combine into a constant in O(1); contiguous
    /// dense buffers combine slice to slice; a dense-and-constant
    /// pair maps the dense operand against the single value; two
    /// dense buffers whose innermost strides are unit or zero
    /// combine run by run as slices; and everything else goes
    /// through the logical-order iterators. A pure `combine` hands
    /// every logical pair the same value on every path, so the
    /// lanes agree bitwise.
    ///
    /// # Panics
    /// Panics if the tensors have different shapes.
    fn zip(&self, other: &Self, combine: impl Fn(&Element, &Element) -> Element) -> Self {
        assert_eq!(
            self.logical_shape(),
            other.logical_shape(),
            "tensors have different shapes"
        );
        if let (Storage::Constant { value: left, .. }, Storage::Constant { value: right, .. }) =
            (&self.storage, &other.storage)
        {
            return Self::constant(self.logical_shape().clone(), combine(left, right));
        }
        if let (Some(left), Some(right)) = (self.as_slice(), other.as_slice()) {
            return Self::dense(
                self.logical_shape().clone(),
                left.iter()
                    .zip(right)
                    .map(|(left, right)| combine(left, right))
                    .collect(),
            );
        }
        if let (Storage::Dense { .. }, Storage::Constant { value: right, .. }) =
            (&self.storage, &other.storage)
        {
            return self.map(|left| combine(left, right));
        }
        if let (Storage::Constant { value: left, .. }, Storage::Dense { .. }) =
            (&self.storage, &other.storage)
        {
            return other.map(|right| combine(left, right));
        }
        if let (
            Storage::Dense {
                data: left,
                layout: left_layout,
            },
            Storage::Dense {
                data: right,
                layout: right_layout,
            },
        ) = (&self.storage, &other.storage)
            && let Some(combined) = zipped_runs(left, left_layout, right, right_layout, &combine)
        {
            return Self::dense(self.logical_shape().clone(), combined);
        }
        Self::dense(
            self.logical_shape().clone(),
            self.iter()
                .zip(other.iter())
                .map(|(left, right)| combine(&left, &right))
                .collect(),
        )
    }
}

/// Combines two same-shape dense buffers by innermost-axis runs, or
/// declines with `None` when either innermost stride exceeds one.
///
/// A unit stride walks a run as a slice and a zero stride holds one
/// element for the whole run, so every accepted case combines slices or
/// a slice against a held element — the loop shapes vectorization needs,
/// where the logical-order odometer defeats it. Both operands advance
/// run by run through their outer-axis offsets in logical order; a
/// zero-by-zero run computes its one value and repeats it.
fn zipped_runs<Element: Clone>(
    left: &[Element],
    left_layout: &Layout,
    right: &[Element],
    right_layout: &Layout,
    combine: impl Fn(&Element, &Element) -> Element,
) -> Option<Vec<Element>> {
    let left_stride = left_layout.inner_stride();
    let right_stride = right_layout.inner_stride();
    if left_stride > 1 || right_stride > 1 {
        return None;
    }
    let extent = left_layout.inner_extent();
    let mut combined = Vec::with_capacity(left_layout.volume());
    for (left_start, right_start) in left_layout.run_offsets().zip(right_layout.run_offsets()) {
        match (left_stride, right_stride) {
            (1, 1) => combined.extend(
                left[left_start..left_start + extent]
                    .iter()
                    .zip(&right[right_start..right_start + extent])
                    .map(|(left, right)| combine(left, right)),
            ),
            (1, 0) => {
                let held = &right[right_start];
                combined.extend(
                    left[left_start..left_start + extent]
                        .iter()
                        .map(|left| combine(left, held)),
                );
            }
            (0, 1) => {
                let held = &left[left_start];
                combined.extend(
                    right[right_start..right_start + extent]
                        .iter()
                        .map(|right| combine(held, right)),
                );
            }
            _ => {
                let value = combine(&left[left_start], &right[right_start]);
                combined.extend(std::iter::repeat_n(value, extent));
            }
        }
    }
    Some(combined)
}

/// Returns the shape with its two axes swapped, matching the payload
/// transpose that a constant undergoes without touching a buffer.
///
/// # Panics
/// Panics if the rank exceeds 2.
fn transpose_shape(shape: &Shape) -> Shape {
    if shape.rank() < 2 {
        return shape.clone();
    }
    assert_eq!(shape.rank(), 2, "transpose supports rank 2 at most");
    let axes = shape.axes();
    Shape::new([axes[1], axes[0]])
}

/// Iterator over a tensor's elements in logical row-major order,
/// yielding owned clones so a representation without a buffer to
/// borrow from can still answer.
///
/// The variants mirror the storage representations: a repeated constant, a
/// direct slice walk for a contiguous buffer, and an odometer walk for a
/// strided view.
enum ElementIter<'tensor, Element> {
    Constant {
        value: &'tensor Element,
        remaining: usize,
    },
    Contiguous(std::slice::Iter<'tensor, Element>),
    Strided {
        data: &'tensor [Element],
        shape: &'tensor [usize],
        strides: &'tensor [usize],
        coordinates: Strides,
        index: usize,
        remaining: usize,
    },
    Selection {
        indices: &'tensor [usize],
        vocab: usize,
        zero: &'tensor Element,
        one: &'tensor Element,
        position: usize,
        total: usize,
    },
}

impl<'tensor, Element: Clone> Iterator for ElementIter<'tensor, Element> {
    type Item = Element;

    fn next(&mut self) -> Option<Element> {
        match self {
            ElementIter::Constant { value, remaining } => {
                if *remaining == 0 {
                    return None;
                }
                *remaining -= 1;
                Some((*value).clone())
            }
            ElementIter::Contiguous(iterator) => iterator.next().cloned(),
            ElementIter::Strided {
                data,
                shape,
                strides,
                coordinates,
                index,
                remaining,
            } => {
                if *remaining == 0 {
                    return None;
                }
                let element = data[*index].clone();
                *remaining -= 1;
                if *remaining > 0 {
                    // Advance the odometer: step the innermost axis, carrying
                    // into the outer axes and adjusting the flat index by the
                    // stride of whichever axis moved.
                    for axis in (0..shape.len()).rev() {
                        coordinates[axis] += 1;
                        if coordinates[axis] < shape[axis] {
                            *index += strides[axis];
                            break;
                        }
                        *index -= (shape[axis] - 1) * strides[axis];
                        coordinates[axis] = 0;
                    }
                }
                Some(element)
            }
            ElementIter::Selection {
                indices,
                vocab,
                zero,
                one,
                position,
                total,
            } => {
                if *position >= *total {
                    return None;
                }
                let row = *position / *vocab;
                let column = *position % *vocab;
                *position += 1;
                Some(if indices[row] == column {
                    (*one).clone()
                } else {
                    (*zero).clone()
                })
            }
        }
    }
}

/// The rank-0 conversion: one element becomes the tensor of shape
/// `[]`. It is what lets `tape.parameter(0.0_f64)` and payload
/// literals in operator position stay scalar-looking while the graph
/// is always tensors. The bound is the `Element` seam marker rather
/// than `Differentiable` so a tensor can never be mistaken for an
/// element of a deeper tensor by inference.
impl<E: Element> From<E> for Tensor<E> {
    fn from(element: E) -> Self {
        Self::constant(Shape::scalar(), element)
    }
}

/// Renders rank 0 as the bare element and higher ranks as nested
/// row-major bracket lists, so a scalar read prints as the number it
/// is.
impl<Element: Clone + fmt::Display> fmt::Display for Tensor<Element> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let elements = self.to_vec();
        display_level(formatter, self.logical_shape().axes(), &elements)
    }
}

/// Writes one nesting level of a tensor display: the bare element at
/// rank 0, a bracketed list of the sub-levels otherwise.
fn display_level<Element: fmt::Display>(
    formatter: &mut fmt::Formatter<'_>,
    axes: &[usize],
    elements: &[Element],
) -> fmt::Result {
    let Some((&extent, rest)) = axes.split_first() else {
        return write!(formatter, "{}", elements[0]);
    };
    let stride = elements.len() / extent;
    write!(formatter, "[")?;
    for index in 0..extent {
        if index > 0 {
            write!(formatter, ", ")?;
        }
        display_level(
            formatter,
            rest,
            &elements[index * stride..(index + 1) * stride],
        )?;
    }
    write!(formatter, "]")
}

impl<Element: PartialEq + Clone> PartialEq for Tensor<Element> {
    /// Compares two tensors by logical value: equal shapes and equal
    /// elements in logical order, independent of storage representation, so
    /// a view compares equal to its materialized twin.
    fn eq(&self, other: &Self) -> bool {
        self.logical_shape() == other.logical_shape() && self.iter().eq(other.iter())
    }
}

impl<Element: Differentiable> Add for Tensor<Element> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        self.zip(&rhs, |left, right| left.clone() + right.clone())
    }
}

impl<Element: Differentiable> Sub for Tensor<Element> {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        self.zip(&rhs, |left, right| left.clone() - right.clone())
    }
}

impl<Element: Differentiable> Mul for Tensor<Element> {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        self.zip(&rhs, |left, right| left.clone() * right.clone())
    }
}

impl<Element: Differentiable> Div for Tensor<Element> {
    type Output = Self;

    fn div(self, rhs: Self) -> Self {
        self.zip(&rhs, |left, right| left.clone() / right.clone())
    }
}

impl<Element: Differentiable> Neg for Tensor<Element> {
    type Output = Self;

    fn neg(self) -> Self {
        self.map(|element| -element.clone())
    }
}

impl<Element: Differentiable> Tensor<Element> {
    /// Returns a zero shaped like `self`, stored as a constant: the
    /// seed of gradient accumulators.
    pub fn zero_like(&self) -> Self {
        Self::constant(self.logical_shape().clone(), Element::zero())
    }

    /// Returns a one shaped like `self`, stored as a constant: the
    /// seed of the output gradient.
    pub fn one_like(&self) -> Self {
        Self::constant(self.logical_shape().clone(), Element::one())
    }

    /// Returns the count spread across `shape`, stored as a constant;
    /// the element value comes from the element type's own
    /// [`from_count`](Differentiable::from_count).
    ///
    /// It is the constructor behind size-derived constants: a composed
    /// formula that divides by an axis extent (a mean, a
    /// normalization) mints that extent here. Counts convert exactly
    /// as long as the element type can represent them.
    ///
    /// # Panics
    /// Panics if the shape's volume overflows `usize` or the shape holds
    /// no elements, as documented on [`Tensor::new`].
    pub fn counted(shape: Shape, count: usize) -> Self {
        Self::constant(shape, Element::from_count(count))
    }

    /// Returns whether this tensor is exactly what
    /// [`counted`](Tensor::counted) mints for `shape` and `count`:
    /// the recognizer pattern matchers use to certify a recorded
    /// size-derived constant before raising a formula around it.
    pub fn is_counted(&self, shape: &Shape, count: usize) -> bool {
        *self.logical_shape() == *shape && self.iter().all(|element| element.is_count(count))
    }
}

impl<Element: Elementary> Tensor<Element> {
    /// Applies one elementwise transcendental: the backend seam first
    /// for a contiguous dense buffer or for a strided view's buffer
    /// window — which keeps the view, so a broadcast operand hands the
    /// backend only its distinct elements — and the scalar `fallback`
    /// everywhere else (constants, declined maps).
    fn mapped(&self, operation: MapOperation, fallback: impl Fn(&Element) -> Element) -> Self {
        if let Some(elements) = self.as_slice()
            && let Some(mapped) = Element::map(&MapTask::new(operation, elements))
        {
            assert_eq!(
                mapped.len(),
                elements.len(),
                "the `Elementary::map` contract requires one output per input element"
            );
            return Self::dense(self.logical_shape().clone(), mapped);
        }
        if let Some((window, layout)) = self.strided_window()
            && let Some(mapped) = Element::map(&MapTask::new(operation, window))
        {
            assert_eq!(
                mapped.len(),
                window.len(),
                "the `Elementary::map` contract requires one output per input element"
            );
            return Self {
                storage: Storage::Dense {
                    data: Arc::new(mapped),
                    layout,
                },
            };
        }
        self.map(fallback)
    }
}

impl<Element: Elementary> Tensor<Element> {
    /// Returns `e` raised to each element.
    pub fn exp(&self) -> Self {
        self.mapped(MapOperation::Exp, |element| element.exp())
    }

    /// Returns the natural logarithm of each element.
    pub fn ln(&self) -> Self {
        self.mapped(MapOperation::Ln, |element| element.ln())
    }

    /// Returns the square root of each element.
    pub fn sqrt(&self) -> Self {
        self.mapped(MapOperation::Sqrt, |element| element.sqrt())
    }

    /// Returns the hyperbolic tangent of each element.
    pub fn tanh(&self) -> Self {
        self.mapped(MapOperation::Tanh, |element| element.tanh())
    }

    /// Returns the sine of each element.
    pub fn sin(&self) -> Self {
        self.mapped(MapOperation::Sin, |element| element.sin())
    }

    /// Returns the cosine of each element.
    pub fn cos(&self) -> Self {
        self.mapped(MapOperation::Cos, |element| element.cos())
    }

    /// Returns the natural logarithm of one plus each element,
    /// accurate near zero.
    pub fn log1p(&self) -> Self {
        self.mapped(MapOperation::Log1p, |element| element.log1p())
    }

    /// Returns `e` raised to each element, minus one, accurate near
    /// zero.
    pub fn expm1(&self) -> Self {
        self.mapped(MapOperation::Expm1, |element| element.expm1())
    }

    /// Returns the error function of each element.
    pub fn erf(&self) -> Self {
        self.mapped(MapOperation::Erf, |element| element.erf())
    }

    /// Returns the derivative of the error function of each element:
    /// the scaled Gaussian `(2/sqrt(pi)) * e^(-x^2)`.
    pub fn erf_derivative(&self) -> Self {
        self.mapped(MapOperation::ErfDerivative, |element| {
            element.erf_derivative()
        })
    }

    /// Returns each element raised to the matching element of
    /// `exponent`.
    ///
    /// # Panics
    /// Panics if the tensors have different shapes.
    pub fn powf(&self, exponent: Self) -> Self {
        self.zip(&exponent, |element, exponent| {
            element.powf(exponent.clone())
        })
    }

    /// Returns the elementwise maximum of `self` and `other`.
    ///
    /// # Panics
    /// Panics if the tensors have different shapes.
    pub fn maximum(&self, other: &Self) -> Self {
        self.zip(other, |element, other| element.maximum(other))
    }

    /// Returns the elementwise 0/1 indicator of `self >= threshold`:
    /// the Heaviside step, ties answering one.
    ///
    /// # Panics
    /// Panics if the tensors have different shapes.
    pub fn step(&self, threshold: &Self) -> Self {
        self.zip(threshold, |element, threshold| element.step(threshold))
    }
}

impl<Element: Elementary> Tensor<Element> {
    /// Returns the batch normalization through the payload seam: when
    /// every operand is a contiguous dense buffer the whole group is
    /// offered to the backend chain as one task, and the composed
    /// bitwise reference computes when the chain declines or an
    /// operand is a view or constant.
    ///
    /// # Panics
    /// Panics if `self` is not rank 2 `[batch, features]`, the affine
    /// operands do not hold `features` values, or `epsilon` holds
    /// more than one value.
    pub fn batch_normalized(
        &self,
        scale: &Self,
        shift: &Self,
        epsilon: &Self,
    ) -> (Self, Self, Self) {
        let shape = self.logical_shape().clone();
        let axes = shape.axes();
        assert_eq!(
            axes.len(),
            2,
            "batch_normalized input must be rank 2 [batch, features], got {shape}"
        );
        let (batch, features) = (axes[0], axes[1]);
        assert_eq!(
            scale.logical_shape().volume(),
            features,
            "batch_normalized scale must hold {features} features"
        );
        assert_eq!(
            shift.logical_shape().volume(),
            features,
            "batch_normalized shift must hold {features} features"
        );
        assert_eq!(
            epsilon.logical_shape().volume(),
            1,
            "batch_normalized epsilon must hold a single value"
        );
        // Only the input must already be contiguous; the affine
        // operands and the epsilon are `features`-sized at most, so
        // materializing them (a constant `filled` scale, a strided
        // view) costs nothing against the fused pass.
        if let Some(input) = self.as_slice() {
            let scale_elements = scale.to_vec();
            let shift_elements = shift.to_vec();
            let epsilon_value = epsilon.iter().next().expect("epsilon holds one value");
            if let Some(normalized) = Element::batch_norm(&BatchNormTask::new(
                input,
                &scale_elements,
                &shift_elements,
                epsilon_value,
                batch,
                features,
            )) {
                let feature_shape = Shape::new([features]);
                return (
                    Self::dense(shape, normalized.output),
                    Self::dense(feature_shape.clone(), normalized.mean),
                    Self::dense(feature_shape, normalized.variance),
                );
            }
        }
        composed_batch_norm(self, scale, shift, epsilon)
    }

    /// Returns the max pool through a direct window walk: each
    /// output element folds its window with `maximum` in the same
    /// row-major lane order as the recorded formula, so the walk is
    /// bit-identical to the composed fold while materializing no
    /// lane views. Non-contiguous inputs and non-dense storages take
    /// the composed reference path.
    ///
    /// # Panics
    /// Panics if `self` is not rank 4, `size` or `stride` is zero,
    /// or a window does not fit the spatial extents.
    pub fn max_pooled(&self, size: usize, stride: usize) -> Self {
        let shape = self.logical_shape();
        let axes = shape.axes();
        assert_eq!(
            axes.len(),
            4,
            "max_pooled input must be rank 4 [batch, channels, height, width], got {shape}"
        );
        assert!(
            size > 0 && stride > 0,
            "max_pooled needs positive size and stride"
        );
        let (batch, channels, height, width) = (axes[0], axes[1], axes[2], axes[3]);
        assert!(
            size <= height && size <= width,
            "max_pooled window {size} does not fit {shape}"
        );
        let Some(elements) = self.as_slice() else {
            return composed_max_pool(self, size, stride);
        };
        let out_height = (height - size) / stride + 1;
        let out_width = (width - size) / stride + 1;
        let mut pooled = Vec::with_capacity(batch * channels * out_height * out_width);
        for image in 0..batch {
            for channel in 0..channels {
                let plane = (image * channels + channel) * height;
                for out_y in 0..out_height {
                    for out_x in 0..out_width {
                        let corner = (plane + out_y * stride) * width + out_x * stride;
                        let mut largest = elements[corner].clone();
                        for lane_y in 0..size {
                            let row = corner + lane_y * width;
                            for lane_x in 0..size {
                                if lane_y == 0 && lane_x == 0 {
                                    continue;
                                }
                                largest = largest.maximum(&elements[row + lane_x]);
                            }
                        }
                        pooled.push(largest);
                    }
                }
            }
        }
        Self::dense(Shape::new([batch, channels, out_height, out_width]), pooled)
    }

    /// Returns the matrix product of two rank-2 tensors.
    ///
    /// Dense operands, including strided views (a transposed operand,
    /// most often), multiply on a slice path that reads their buffers
    /// through the layout strides directly; other storages read
    /// through logical access. Both paths accumulate every output
    /// element in the same order, so their results are bit-identical.
    ///
    /// # Panics
    /// Panics if either operand is not rank 2, the inner dimensions do not
    /// agree, or any dimension is empty.
    pub fn matmul(&self, rhs: &Self) -> Self {
        let left = self.logical_shape();
        let right = rhs.logical_shape();
        assert!(left.rank() >= 2, "matmul requires rank-2 or higher tensors");
        assert_eq!(
            left.rank(),
            right.rank(),
            "matmul operands must agree in rank"
        );
        let split = left.rank() - 2;
        assert_eq!(
            &left.axes()[..split],
            &right.axes()[..split],
            "matmul batch axes do not agree"
        );
        let (rows, inner) = (left.axes()[split], left.axes()[split + 1]);
        let (rhs_inner, columns) = (right.axes()[split], right.axes()[split + 1]);
        assert_eq!(inner, rhs_inner, "matmul inner dimensions do not agree");
        assert!(
            rows > 0 && inner > 0 && columns > 0,
            "matmul requires non-empty dimensions"
        );
        let batch_axes = &left.axes()[..split];
        let batches: usize = batch_axes.iter().product();
        let result_shape = Shape::new(batch_axes.iter().copied().chain([rows, columns]));
        if batches == 0 {
            return Self::dense(result_shape, Vec::new());
        }

        if let (Some((a, a_strides)), Some((b, b_strides))) =
            (self.gemm_operand(), rhs.gemm_operand())
        {
            // Each batch slice is itself a valid 2D strided view, so
            // the loop issues one rank-2 task per slice through the
            // same seam — the backends and the accumulator contract
            // are inherited unchanged, bitwise per slice. Tier one is
            // the element's acceleration seam (the compiled backend
            // chain); tier two the built-in slice path. A declined
            // task costs one call answering `None`.
            let mut elements = Vec::with_capacity(batches * rows * columns);
            let mut index = vec![0usize; split];
            loop {
                let a_base: usize = index
                    .iter()
                    .zip(&a_strides[..split])
                    .map(|(&at, &stride)| at * stride)
                    .sum();
                let b_base: usize = index
                    .iter()
                    .zip(&b_strides[..split])
                    .map(|(&at, &stride)| at * stride)
                    .sum();
                let task = gemm::GemmTask::new(
                    &a[a_base..],
                    [a_strides[split], a_strides[split + 1]],
                    &b[b_base..],
                    [b_strides[split], b_strides[split + 1]],
                    rows,
                    inner,
                    columns,
                );
                match Element::gemm(&task) {
                    Some(product) => {
                        assert_eq!(
                            product.len(),
                            rows * columns,
                            "the `Elementary::gemm` contract requires `rows * columns` elements"
                        );
                        elements.extend(product);
                    }
                    None => elements.extend(gemm::multiply(&task)),
                }
                // The batch odometer over the leading axes.
                let mut axis = split;
                loop {
                    if axis == 0 {
                        return Self::dense(result_shape, elements);
                    }
                    axis -= 1;
                    index[axis] += 1;
                    if index[axis] < batch_axes[axis] {
                        break;
                    }
                    index[axis] = 0;
                }
            }
        }

        let mut elements = Vec::with_capacity(batches * rows * columns);
        for batch in 0..batches {
            let a_base = batch * rows * inner;
            let b_base = batch * inner * columns;
            for row in 0..rows {
                for column in 0..columns {
                    let mut total = self.get(a_base + row * inner).promote()
                        * rhs.get(b_base + column).promote();
                    for step in 1..inner {
                        total = total
                            + self.get(a_base + row * inner + step).promote()
                                * rhs.get(b_base + step * columns + column).promote();
                    }
                    elements.push(Element::demote(total));
                }
            }
        }
        Self::dense(result_shape, elements)
    }

    /// Returns the tensor with its two axes swapped as a view over the same
    /// buffer.
    ///
    /// Rank-0 and rank-1 tensors are returned unchanged.
    ///
    /// # Panics
    /// Panics if the tensor's rank exceeds 2.
    pub fn transpose(&self) -> Self {
        match &self.storage {
            Storage::Dense { data, layout } => Self {
                storage: Storage::Dense {
                    data: Arc::clone(data),
                    layout: layout.transpose(),
                },
            },
            Storage::Constant { shape, value } => {
                Self::constant(transpose_shape(shape), value.clone())
            }
            Storage::Selection { .. } => self.densify().transpose(),
        }
    }

    /// Returns the sum of every element as a rank-0 constant.
    ///
    /// Elements are accumulated in logical order from left to right without
    /// pairwise or compensated summation.
    pub fn sum(&self) -> Self {
        let mut elements = self.iter();
        let first = elements
            .next()
            .expect("sum requires a non-empty tensor")
            .promote();
        let total = elements.fold(first, |total, element| total + element.promote());
        Self::constant(Shape::scalar(), Element::demote(total))
    }

    /// Returns the tensor with `axis` reduced by summation.
    ///
    /// The reduction is rank-general: the elements are viewed as
    /// `[outer, axis, inner]` in logical order and summed over the middle
    /// extent.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn sum_along(&self, axis: usize) -> Self {
        let shape = self.logical_shape();
        let axes = shape.axes();
        assert!(axis < axes.len(), "axis {axis} is out of rank for {shape}");
        let outer: usize = axes[..axis].iter().product();
        let extent = axes[axis];
        let inner: usize = axes[axis + 1..].iter().product();

        let mut elements = Vec::with_capacity(outer * inner);
        for outer_index in 0..outer {
            for inner_index in 0..inner {
                let position = |step: usize| (outer_index * extent + step) * inner + inner_index;
                let mut total = self.get(position(0)).promote();
                for step in 1..extent {
                    total = total + self.get(position(step)).promote();
                }
                elements.push(Element::demote(total));
            }
        }
        Self::dense(shape.without_axis(axis), elements)
    }

    /// Returns `self` with `axis` reduced by the stable log-sum-exp:
    /// the axis maximum plus the log of the shifted exponential sum.
    ///
    /// Shifting by the axis maximum keeps every exponent at or below
    /// zero: the sum lands between one and the axis extent, its
    /// logarithm between zero and `ln(extent)`, so the result is
    /// finite for every finite operand — even where the shifted
    /// difference itself underflows to `-inf`.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn logsumexp(&self, axis: usize) -> Self {
        let peak = self.max_along(axis);
        let shifted = self.clone() - peak.broadcast_along_like(axis, self);
        peak + shifted.exp().sum_along(axis).ln()
    }

    /// Returns the log-softmax of `self` along `axis`:
    /// `x - ln(sum(exp(x)))`, the logarithm of the softmax
    /// probabilities.
    ///
    /// Shifting by the axis maximum keeps every exponent at or below
    /// zero, so the sum cannot overflow; the shift cancels in the
    /// final subtraction, leaving the result stable (not exact: the
    /// shifted rounding differs from the unshifted ideal, and a
    /// difference beyond the representable range still underflows to
    /// `-inf` — the mathematically faithful log-probability).
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn log_softmax(&self, axis: usize) -> Self {
        let peak = self.max_along(axis).broadcast_along_like(axis, self);
        let shifted = self.clone() - peak;
        let normalizer = shifted.exp().sum_along(axis).ln();
        shifted.clone() - normalizer.broadcast_along_like(axis, &shifted)
    }

    /// Returns the tensor with `axis` reduced to its largest element by the
    /// elementwise [`Elementary::maximum`].
    ///
    /// The reduction is rank-general and mirrors [`Recordable::sum_along`]:
    /// the elements are viewed as `[outer, axis, inner]` in logical order
    /// and folded over the middle extent.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank.
    pub fn max_along(&self, axis: usize) -> Self {
        let shape = self.logical_shape();
        let axes = shape.axes();
        assert!(axis < axes.len(), "axis {axis} is out of rank for {shape}");
        let outer: usize = axes[..axis].iter().product();
        let extent = axes[axis];
        let inner: usize = axes[axis + 1..].iter().product();

        let mut elements = Vec::with_capacity(outer * inner);
        for outer_index in 0..outer {
            for inner_index in 0..inner {
                let position = |step: usize| (outer_index * extent + step) * inner + inner_index;
                let mut largest = self.get(position(0));
                for step in 1..extent {
                    largest = largest.maximum(&self.get(position(step)));
                }
                elements.push(largest);
            }
        }
        Self::dense(shape.without_axis(axis), elements)
    }

    /// Returns this tensor's single element spread across `shape` as a
    /// constant: the whole-shape form of explicit broadcasting.
    ///
    /// # Panics
    /// Panics if `self` holds more than one element.
    pub fn broadcast(&self, shape: Shape) -> Self {
        assert_eq!(
            self.logical_shape().volume(),
            1,
            "broadcast requires a single-element tensor"
        );
        Self::constant(shape, self.get(0))
    }

    /// Returns this tensor's single element spread across `reference`'s
    /// shape: [`broadcast`](Self::broadcast) reading the reference for
    /// its shape alone.
    ///
    /// # Panics
    /// Panics if `self` holds more than one element.
    pub fn broadcast_like(&self, reference: &Self) -> Self {
        self.broadcast(reference.logical_shape().clone())
    }

    /// Returns the tensor repeated along a new axis of `extent`
    /// inserted at `axis`, as a stride-0 view: the named-axis form of
    /// explicit broadcasting.
    ///
    /// # Panics
    /// Panics if `axis` exceeds this tensor's rank or `extent` is zero.
    pub fn broadcast_along(&self, axis: usize, extent: usize) -> Self {
        let shape = self.logical_shape();
        assert!(
            axis <= shape.rank(),
            "broadcast axis {axis} is out of rank for {shape}"
        );
        assert!(extent > 0, "broadcast extent must be positive");
        let mut axes: Vec<usize> = shape.axes().to_vec();
        axes.insert(axis, extent);
        let widened = Shape::new(axes);
        match &self.storage {
            Storage::Constant { value, .. } => Self::constant(widened, value.clone()),
            Storage::Dense { data, layout } => Self {
                storage: Storage::Dense {
                    data: Arc::clone(data),
                    layout: layout.broadcast_along(axis, &widened),
                },
            },
            Storage::Selection { .. } => self.densify().broadcast_along(axis, extent),
        }
    }

    /// Returns the tensor repeated along `axis` to match `reference`'s
    /// shape: [`broadcast_along`](Self::broadcast_along) reading the
    /// reference for its shape alone.
    ///
    /// # Panics
    /// Panics if `axis` is out of `reference`'s rank or `self`'s shape
    /// differs from `reference`'s with that axis removed.
    pub fn broadcast_along_like(&self, axis: usize, reference: &Self) -> Self {
        let reference_shape = reference.logical_shape();
        assert!(
            axis < reference_shape.rank(),
            "axis {axis} is out of rank for {reference_shape}"
        );
        assert_eq!(
            self.logical_shape(),
            &reference_shape.without_axis(axis),
            "broadcast along axis {axis} of {reference_shape} requires the remaining shape"
        );
        self.broadcast_along(axis, reference_shape.axes()[axis])
    }

    /// Returns `self` reinterpreted with `shape` in logical row-major
    /// order.
    ///
    /// A contiguous dense tensor and a constant reshape into an O(1) view
    /// over the same buffer, and so does a strided view when only extent-1
    /// axes are inserted or removed; any other strided reshape is first
    /// materialized.
    ///
    /// # Panics
    /// Panics if `shape`'s volume differs from `self`'s.
    pub fn reshape(&self, shape: Shape) -> Self {
        assert_eq!(
            self.logical_shape().volume(),
            shape.volume(),
            "reshape from {} to {shape} changes the number of elements",
            self.logical_shape()
        );
        match &self.storage {
            Storage::Constant { value, .. } => Self::constant(shape, value.clone()),
            Storage::Dense { data, layout } => match layout.reshape(shape.clone()) {
                Some(reshaped) => Self {
                    storage: Storage::Dense {
                        data: Arc::clone(data),
                        layout: reshaped,
                    },
                },
                None => Self::dense(shape, self.to_vec()),
            },
            Storage::Selection { .. } => Self::dense(shape, self.to_vec()),
        }
    }

    /// Returns `self` with its axes reordered by `order` as a view over the
    /// same buffer.
    ///
    /// # Panics
    /// Panics if `order` is not a permutation of `0..rank`.
    pub fn permute(&self, order: &[usize]) -> Self {
        let shape = self.logical_shape();
        assert_eq!(
            order.len(),
            shape.rank(),
            "permute order must cover every axis of {shape}"
        );
        let mut seen = vec![false; shape.rank()];
        for &axis in order {
            assert!(
                axis < shape.rank(),
                "permute axis {axis} is out of rank for {shape}"
            );
            assert!(
                !std::mem::replace(&mut seen[axis], true),
                "permute order repeats axis {axis}"
            );
        }
        match &self.storage {
            Storage::Constant { value, .. } => {
                let axes = shape.axes();
                let permuted = Shape::new(order.iter().map(|&axis| axes[axis]));
                Self::constant(permuted, value.clone())
            }
            Storage::Dense { data, layout } => Self {
                storage: Storage::Dense {
                    data: Arc::clone(data),
                    layout: layout.permute(order),
                },
            },
            Storage::Selection { .. } => self.densify().permute(order),
        }
    }

    /// Returns the window of `len` elements from `start` along `axis` as a
    /// view over the same buffer.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank, `len` is zero (tensors cannot be
    /// empty), or `start + len` overflows or exceeds the axis extent.
    pub fn narrow(&self, axis: usize, start: usize, len: usize) -> Self {
        let shape = self.logical_shape();
        assert!(
            axis < shape.rank(),
            "narrow axis {axis} is out of rank for {shape}"
        );
        assert!(len > 0, "narrow window must hold at least one element");
        let extent = shape.axes()[axis];
        let end = start
            .checked_add(len)
            .expect("narrow window end overflows `usize`");
        assert!(
            end <= extent,
            "narrow window {start}..{end} exceeds axis {axis} extent {extent}"
        );
        match &self.storage {
            Storage::Constant { value, .. } => {
                let narrowed = Shape::new(
                    shape
                        .axes()
                        .iter()
                        .enumerate()
                        .map(|(index, &e)| if index == axis { len } else { e }),
                );
                Self::constant(narrowed, value.clone())
            }
            Storage::Dense { data, layout } => Self {
                storage: Storage::Dense {
                    data: Arc::clone(data),
                    layout: layout.narrow(axis, start, len),
                },
            },
            Storage::Selection { .. } => self.densify().narrow(axis, start, len),
        }
    }

    /// Returns `self` placed at `start ..` along `axis` inside a tensor
    /// whose `axis` has extent `full_extent`, with zeros elsewhere.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank or the window overflows or exceeds
    /// `full_extent`.
    pub fn pad(&self, axis: usize, start: usize, full_extent: usize) -> Self {
        let shape = self.logical_shape();
        assert!(
            axis < shape.rank(),
            "pad axis {axis} is out of rank for {shape}"
        );
        let axes = shape.axes();
        let len = axes[axis];
        let end = start
            .checked_add(len)
            .expect("pad window end overflows `usize`");
        assert!(
            end <= full_extent,
            "pad window {start}..{end} exceeds the full extent {full_extent}"
        );
        let outer: usize = axes[..axis].iter().product();
        let inner: usize = axes[axis + 1..].iter().product();
        let zero = Element::zero();

        let mut elements = Vec::with_capacity(outer * full_extent * inner);
        for outer_index in 0..outer {
            for position in 0..full_extent {
                for inner_index in 0..inner {
                    if position >= start && position < end {
                        let source = (outer_index * len + (position - start)) * inner + inner_index;
                        elements.push(self.get(source));
                    } else {
                        elements.push(zero.clone());
                    }
                }
            }
        }
        let padded = Shape::new(
            axes.iter()
                .enumerate()
                .map(|(index, &e)| if index == axis { full_extent } else { e }),
        );
        Self::dense(padded, elements)
    }

    /// Returns the sliding windows of `self` along `axis` as a strided
    /// view over the same buffer: the axis becomes a `(count, size)`
    /// pair where window `w` starts at `w * step` and takes every
    /// `dilation`-th element. Overlapping windows alias elements
    /// read-only, which immutability makes safe.
    ///
    /// # Panics
    /// Panics if `axis` is out of rank, `size`, `step`, or `dilation`
    /// is zero, or the dilated window span `dilation * (size - 1) + 1`
    /// overflows or exceeds the axis extent.
    pub fn unfold(&self, axis: usize, size: usize, step: usize, dilation: usize) -> Self {
        let shape = self.logical_shape();
        assert!(
            axis < shape.rank(),
            "unfold axis {axis} is out of rank for {shape}"
        );
        assert!(size > 0, "unfold windows must hold at least one element");
        assert!(step > 0, "unfold step must be positive");
        assert!(dilation > 0, "unfold dilation must be positive");
        let extent = shape.axes()[axis];
        let span = dilation
            .checked_mul(size - 1)
            .and_then(|reach| reach.checked_add(1))
            .expect("unfold window span overflows `usize`");
        assert!(
            span <= extent,
            "unfold window span {span} exceeds axis {axis} extent {extent}"
        );
        match &self.storage {
            Storage::Constant { value, .. } => {
                let count = (extent - span) / step + 1;
                let mut unfolded: Vec<usize> = shape.axes().to_vec();
                unfolded[axis] = count;
                unfolded.insert(axis + 1, size);
                Self::constant(Shape::new(unfolded), value.clone())
            }
            Storage::Dense { data, layout } => Self {
                storage: Storage::Dense {
                    data: Arc::clone(data),
                    layout: layout.unfold(axis, size, step, dilation),
                },
            },
            Storage::Selection { .. } => self.densify().unfold(axis, size, step, dilation),
        }
    }

    /// Returns the `(count, size)` window pair at `axis`, `axis + 1`
    /// folded back onto an axis of `extent`: the adjoint of
    /// [`unfold`](Recordable::unfold) and its gradient rule.
    ///
    /// The accumulation is output-centric: each source position sums,
    /// in window order, the window elements that were read from it, so
    /// the result is deterministic under any evaluation strategy.
    /// Positions no window reaches fold to zero.
    ///
    /// # Panics
    /// Panics if `axis + 1` is out of rank, `size`, `step`, or
    /// `dilation` is zero, the dilated window span overflows or exceeds
    /// `extent`, or the shape at `axis`, `axis + 1` disagrees with the
    /// windows `unfold` would produce for `extent`.
    pub fn fold(
        &self,
        axis: usize,
        size: usize,
        step: usize,
        dilation: usize,
        extent: usize,
    ) -> Self {
        let shape = self.logical_shape();
        let axes = shape.axes();
        assert!(
            axis + 1 < axes.len(),
            "fold window axes {axis}, {} are out of rank for {shape}",
            axis + 1
        );
        assert!(size > 0, "fold windows must hold at least one element");
        assert!(step > 0, "fold step must be positive");
        assert!(dilation > 0, "fold dilation must be positive");
        let span = dilation
            .checked_mul(size - 1)
            .and_then(|reach| reach.checked_add(1))
            .expect("fold window span overflows `usize`");
        assert!(
            span <= extent,
            "fold window span {span} exceeds the extent {extent}"
        );
        let count = (extent - span) / step + 1;
        assert_eq!(
            axes[axis], count,
            "fold window count {} disagrees with the {count} windows of extent {extent}",
            axes[axis]
        );
        assert_eq!(
            axes[axis + 1],
            size,
            "fold window size {} disagrees with {size}",
            axes[axis + 1]
        );
        let outer: usize = axes[..axis].iter().product();
        let inner: usize = axes[axis + 2..].iter().product();
        let zero = Element::zero();

        let mut elements = Vec::with_capacity(outer * extent * inner);
        for outer_index in 0..outer {
            for position in 0..extent {
                for inner_index in 0..inner {
                    let mut total = zero.promote();
                    // Window `w` reads `position` as its element `k`
                    // exactly when `w * step + k * dilation == position`.
                    for k in 0..size {
                        let reach = k * dilation;
                        if reach > position {
                            break;
                        }
                        let rest = position - reach;
                        if !rest.is_multiple_of(step) {
                            continue;
                        }
                        let window = rest / step;
                        if window >= count {
                            continue;
                        }
                        let source =
                            ((outer_index * count + window) * size + k) * inner + inner_index;
                        total = total + self.get(source).promote();
                    }
                    elements.push(Element::demote(total));
                }
            }
        }
        let folded = Shape::new(
            axes[..axis]
                .iter()
                .copied()
                .chain(std::iter::once(extent))
                .chain(axes[axis + 2..].iter().copied()),
        );
        Self::dense(folded, elements)
    }

    /// Returns the im2col matrix through a specialized patch fill:
    /// contiguous runs of `kernel_width` copied per channel and kernel
    /// row, zero runs where the padding window leaves the input, and no
    /// per-element odometer arithmetic — the measured cost the fused
    /// window-GEMM pattern exists to remove. Non-contiguous inputs and
    /// non-dense storages take the composed reference path.
    ///
    /// # Panics
    /// Panics if `self` is not rank 4, `stride` is zero, or a kernel
    /// window does not fit the padded extents.
    pub fn windowed_patches(
        &self,
        kernel_height: usize,
        kernel_width: usize,
        stride: usize,
        padding: usize,
    ) -> Self {
        let shape = self.logical_shape();
        let axes = shape.axes();
        assert_eq!(
            axes.len(),
            4,
            "windowed_product input must be rank 4 [batch, channels, height, width], got {shape}"
        );
        assert!(stride > 0, "windowed_product stride must be positive");
        let (batch, channels, height, width) = (axes[0], axes[1], axes[2], axes[3]);
        assert!(
            kernel_height > 0
                && kernel_width > 0
                && kernel_height <= height + 2 * padding
                && kernel_width <= width + 2 * padding,
            "windowed_product kernel {kernel_height}x{kernel_width} does not fit {shape} \
             with padding {padding}"
        );
        let Some(elements) = self.as_slice() else {
            return composed_windowed_patches(self, kernel_height, kernel_width, stride, padding);
        };

        let out_height = (height + 2 * padding - kernel_height) / stride + 1;
        let out_width = (width + 2 * padding - kernel_width) / stride + 1;
        let columns = channels * kernel_height * kernel_width;
        let zero = Element::zero();
        let mut patches = vec![zero; batch * out_height * out_width * columns];
        for image in 0..batch {
            for out_y in 0..out_height {
                for out_x in 0..out_width {
                    let row = ((image * out_height + out_y) * out_width + out_x) * columns;
                    let source_x = (out_x * stride) as isize - padding as isize;
                    // The kernel columns that land inside the image; the
                    // rest of the patch row stays zero (the padding).
                    let clip_low = (-source_x).max(0) as usize;
                    let clip_high = kernel_width.min((width as isize - source_x).max(0) as usize);
                    if clip_low >= clip_high {
                        continue;
                    }
                    let run = clip_high - clip_low;
                    for channel in 0..channels {
                        for kernel_y in 0..kernel_height {
                            let source_y = (out_y * stride + kernel_y) as isize - padding as isize;
                            if source_y < 0 || source_y >= height as isize {
                                continue;
                            }
                            let source =
                                ((image * channels + channel) * height + source_y as usize) * width
                                    + (source_x + clip_low as isize) as usize;
                            let target = row
                                + (channel * kernel_height + kernel_y) * kernel_width
                                + clip_low;
                            patches[target..target + run]
                                .clone_from_slice(&elements[source..source + run]);
                        }
                    }
                }
            }
        }
        Self::dense(
            Shape::new([batch * out_height * out_width, columns]),
            patches,
        )
    }

    /// Returns the rows of `self` selected by `selection`, a one-hot
    /// `[count, vocab]` whose vocabulary must equal `self`'s first axis; the
    /// result is `[count, ...self.shape[1..]]` with row `i` equal to
    /// `self`'s row `selection_index(i)`.
    ///
    /// # Panics
    /// Panics if `selection` is not a `[count, vocab]` selection, `self` has
    /// no axes, or the vocabulary does not match `self`'s first axis.
    pub fn gather(&self, selection: &Self) -> Self {
        let table = self.logical_shape();
        let indices = selection.selection_indices();
        assert!(table.rank() >= 1, "gather table needs at least one axis");
        let vocabulary = selection.logical_shape().axes()[1];
        assert_eq!(
            vocabulary,
            table.axes()[0],
            "gather selection vocabulary {vocabulary} does not match table rows {}",
            table.axes()[0]
        );
        let row_size: usize = table.axes()[1..].iter().product();

        let mut elements = Vec::with_capacity(indices.len() * row_size);
        for &row in indices {
            for offset in 0..row_size {
                elements.push(self.get(row * row_size + offset));
            }
        }
        let result =
            Shape::new(std::iter::once(indices.len()).chain(table.axes()[1..].iter().copied()));
        Self::dense(result, elements)
    }

    /// Scatter-adds the rows of `self` (a `[count, ...]` gradient) into a
    /// zero payload with one row per entry of `selection`'s vocabulary,
    /// by its indices: the adjoint of [`gather`](Recordable::gather) and
    /// its gradient rule. Rows selected more than once accumulate.
    pub fn scatter(&self, selection: &Self) -> Self {
        let gradient = self.logical_shape();
        assert!(
            gradient.rank() >= 1,
            "scatter needs a gradient with a leading selection axis"
        );
        let indices = selection.selection_indices();
        assert_eq!(
            gradient.axes()[0],
            indices.len(),
            "scatter gradient rows disagree with the selection count"
        );
        let rows = selection.logical_shape().axes()[1];
        let row_size: usize = gradient.axes()[1..].iter().product();
        let volume = rows
            .checked_mul(row_size)
            .expect("shape volume overflows `usize`");
        let zero = Element::zero();

        let mut accumulators = vec![zero.promote(); volume];
        for (source, &target) in indices.iter().enumerate() {
            for offset in 0..row_size {
                let position = target * row_size + offset;
                accumulators[position] =
                    accumulators[position].clone() + self.get(source * row_size + offset).promote();
            }
        }
        let result = Shape::new(std::iter::once(rows).chain(gradient.axes()[1..].iter().copied()));
        Self::dense(
            result,
            accumulators.into_iter().map(Element::demote).collect(),
        )
    }
}

impl<Element: Elementary> Tensor<Element> {
    /// Returns the im2col product of `self` with the GEMM-shaped
    /// `kernel`: the window rows of the padded, strided sliding
    /// windows, matrix-multiplied in one call. It is the fused
    /// executor behind the plan tier's window-GEMM pattern; the
    /// composed reference is
    /// [`composed_windowed_patches`](crate::reference::composed_windowed_patches)
    /// followed by the plain product.
    pub fn windowed_product(
        &self,
        kernel: &Self,
        kernel_height: usize,
        kernel_width: usize,
        stride: usize,
        padding: usize,
    ) -> Self {
        self.windowed_patches(kernel_height, kernel_width, stride, padding)
            .matmul(kernel)
    }
}

/// The compute interpretation of the recordable vocabulary: every
/// rule operation is the inherent tensor method of the same name, so
/// running a derivative rule over tensors is running it over the
/// engine's one payload.
impl<E: Element> Recordable for Tensor<E> {
    fn shape(&self) -> Shape {
        Tensor::shape(self)
    }

    fn zero_like(&self) -> Self {
        Tensor::zero_like(self)
    }

    fn one_like(&self) -> Self {
        Tensor::one_like(self)
    }

    fn exp(&self) -> Self {
        Tensor::exp(self)
    }

    fn ln(&self) -> Self {
        Tensor::ln(self)
    }

    fn sqrt(&self) -> Self {
        Tensor::sqrt(self)
    }

    fn tanh(&self) -> Self {
        Tensor::tanh(self)
    }

    fn sin(&self) -> Self {
        Tensor::sin(self)
    }

    fn cos(&self) -> Self {
        Tensor::cos(self)
    }

    fn log1p(&self) -> Self {
        Tensor::log1p(self)
    }

    fn expm1(&self) -> Self {
        Tensor::expm1(self)
    }

    fn erf(&self) -> Self {
        Tensor::erf(self)
    }

    fn erf_derivative(&self) -> Self {
        Tensor::erf_derivative(self)
    }

    fn powf(&self, exponent: Self) -> Self {
        Tensor::powf(self, exponent)
    }

    fn maximum(&self, other: &Self) -> Self {
        Tensor::maximum(self, other)
    }

    fn step(&self, threshold: &Self) -> Self {
        Tensor::step(self, threshold)
    }

    fn matmul(&self, rhs: &Self) -> Self {
        Tensor::matmul(self, rhs)
    }

    fn sum(&self) -> Self {
        Tensor::sum(self)
    }

    fn sum_along(&self, axis: usize) -> Self {
        Tensor::sum_along(self, axis)
    }

    fn logsumexp(&self, axis: usize) -> Self {
        Tensor::logsumexp(self, axis)
    }

    fn log_softmax(&self, axis: usize) -> Self {
        Tensor::log_softmax(self, axis)
    }

    fn broadcast(&self, shape: Shape) -> Self {
        Tensor::broadcast(self, shape)
    }

    fn broadcast_along(&self, axis: usize, extent: usize) -> Self {
        Tensor::broadcast_along(self, axis, extent)
    }

    fn reshape(&self, shape: Shape) -> Self {
        Tensor::reshape(self, shape)
    }

    fn permute(&self, order: &[usize]) -> Self {
        Tensor::permute(self, order)
    }

    fn narrow(&self, axis: usize, start: usize, len: usize) -> Self {
        Tensor::narrow(self, axis, start, len)
    }

    fn pad(&self, axis: usize, start: usize, full_extent: usize) -> Self {
        Tensor::pad(self, axis, start, full_extent)
    }

    fn unfold(&self, axis: usize, size: usize, step: usize, dilation: usize) -> Self {
        Tensor::unfold(self, axis, size, step, dilation)
    }

    fn fold(&self, axis: usize, size: usize, step: usize, dilation: usize, extent: usize) -> Self {
        Tensor::fold(self, axis, size, step, dilation, extent)
    }

    fn gather(&self, selection: &Self) -> Self {
        Tensor::gather(self, selection)
    }

    fn scatter(&self, selection: &Self) -> Self {
        Tensor::scatter(self, selection)
    }
}

#[cfg(test)]
#[path = "tests/tensor_tests.rs"]
mod tests;