nuvec-core 0.6.0

//! Batching rules: `Primitive::vmap`.
//!
//! Each rule receives the real batched inputs with their batch axes (`None` for unbatched
//! inputs) and records the batched computation. Rules that change axes first move the batch
//! axis to the front; element-wise rules without axis changes keep it in place. Every
//! primitive has a rule (a missing one fails to compile), so `vmap` needs no loop fallback.

use crate::array::Array;
use crate::ops::imp::{broadcast_to, expand_dims, matmul, reshape, transpose};
use crate::primitive::Primitive;
use crate::quantized::{dequantize, quantize, quantized_matmul};
use crate::{Result, view};

fn record(primitive: Primitive, inputs: Vec<Array>) -> Result<Vec<Array>> {
    Array::from_op(primitive, inputs)
}

/// `a` with axis `from` moved to position `to`.
pub(crate) fn move_axis(a: &Array, from: usize, to: usize) -> Result<Array> {
    if from == to {
        return Ok(a.clone());
    }
    let mut perm: Vec<usize> = (0..a.ndim()).filter(|&i| i != from).collect();
    perm.insert(to, from);
    transpose(a, &perm)
}

/// `a` with its batch axis first; unbatched inputs get a batch axis of size 1.
fn front(a: &Array, axis: Option<usize>) -> Result<Array> {
    match axis {
        Some(axis) => move_axis(a, axis, 0),
        None => expand_dims(a, 0),
    }
}

/// [`front`], with unbatched inputs broadcast to the full batch size `b`.
fn front_full(a: &Array, axis: Option<usize>, b: usize) -> Result<Array> {
    let x = front(a, axis)?;
    let mut shape = x.dims().to_vec();
    shape[0] = b;
    broadcast_to(&x, &shape)
}

/// The batch axis of a gathered input that `vmap` moved last (a batched vector gathers
/// as one row per rank, its batch on the second axis).
fn moved_last(ndim: usize) -> usize {
    ndim.max(2) - 1
}

/// Prepends a batch dimension to slice parameters.
fn prepend<T: Copy>(first: T, rest: &[T]) -> Vec<T> {
    std::iter::once(first).chain(rest.iter().copied()).collect()
}

impl Primitive {
    /// The batched outputs of this primitive for batched `inputs` (batch `axes`, `None`
    /// for unbatched inputs; at least one is batched), and their common batch axis.
    pub(crate) fn vmap(
        &self,
        inputs: &[Array],
        axes: &[Option<usize>],
    ) -> Result<(Vec<Array>, usize)> {
        let b = inputs
            .iter()
            .zip(axes)
            .find_map(|(a, ax)| ax.map(|ax| a.dims()[ax]))
            .expect("at least one batched input");
        let x = &inputs[0];
        let at_front = |p: Primitive| Ok((record(p, vec![front(x, axes[0])?])?, 0));
        match self {
            Self::Depends { .. } => {
                // Batched `x` keeps its axis; the dependencies only order the work.
                let (x, axis) = match axes[0] {
                    Some(axis) => (x.clone(), axis),
                    None => (front_full(x, None, b)?, 0),
                };
                let mut all = vec![x];
                all.extend_from_slice(&inputs[1..]);
                Ok((record(self.clone(), all)?, axis))
            }
            // MLX's rules: the reductions and sends keep the batch axis; a gather joins the
            // ranks along the first axis, so the batch goes last.
            Self::Collective { op, .. } => match op {
                crate::distributed::CollectiveOp::Gather => {
                    let axis = axes[0].expect("the one input is batched");
                    let last = x.ndim() - 1;
                    let moved = if x.ndim() == 1 {
                        crate::ops::imp::expand_dims(x, 0)?
                    } else {
                        move_axis(x, axis, last)?
                    };
                    Ok((record(self.clone(), vec![moved])?, moved_last(x.ndim())))
                }
                crate::distributed::CollectiveOp::Recv { .. } => {
                    unreachable!("recv has no inputs to batch")
                }
                // The scatter runs along the first axis: a batch there moves to the second.
                crate::distributed::CollectiveOp::SumScatter => {
                    let axis = axes[0].expect("the one input is batched");
                    if axis == 0 {
                        let moved = move_axis(x, 0, 1)?;
                        Ok((record(self.clone(), vec![moved])?, 1))
                    } else {
                        Ok((record(self.clone(), vec![x.clone()])?, axis))
                    }
                }
                _ => Ok((record(self.clone(), vec![x.clone()])?, axes[0].unwrap_or(0))),
            },
            Self::Unary(_) | Self::AsType(_) | Self::StopGradient | Self::Copy => {
                Ok((record(self.clone(), vec![x.clone()])?, axes[0].unwrap_or(0)))
            }
            Self::ToFp8 | Self::FromFp8(_) => {
                Ok((record(self.clone(), vec![x.clone()])?, axes[0].unwrap_or(0)))
            }
            // Per row of the last axis: the batch goes first.
            Self::SortedOffsets { .. } => Ok((record(self.clone(), vec![front(x, axes[0])?])?, 0)),
            // The last axis is reinterpreted: the batch goes first.
            Self::View(_) => Ok((record(self.clone(), vec![front(x, axes[0])?])?, 0)),
            Self::Binary(_) | Self::Select => self.elementwise(inputs, axes),
            Self::Reshape(shape) => {
                let shape: Vec<isize> = prepend(b, shape).iter().map(|&n| n as isize).collect();
                Ok((vec![reshape(&front(x, axes[0])?, &shape)?], 0))
            }
            Self::Squeeze(unit) => at_front(Self::Squeeze(unit.iter().map(|a| a + 1).collect())),
            Self::ExpandDims(unit) => {
                at_front(Self::ExpandDims(unit.iter().map(|a| a + 1).collect()))
            }
            Self::Transpose(perm) => {
                let perm: Vec<usize> = prepend(0, &perm.iter().map(|a| a + 1).collect::<Vec<_>>());
                Ok((vec![transpose(&front(x, axes[0])?, &perm)?], 0))
            }
            Self::Slice { start, step, len } => at_front(Self::Slice {
                start: prepend(0, start),
                step: prepend(1, step),
                len: prepend(b, len),
            }),
            Self::SliceUpdate { start, step, len } => {
                let batched = Self::SliceUpdate {
                    start: prepend(0, start),
                    step: prepend(1, step),
                    len: prepend(b, len),
                };
                let base = front_full(&inputs[0], axes[0], b)?;
                let update = front_full(&inputs[1], axes[1], b)?;
                Ok((record(batched, vec![base, update])?, 0))
            }
            Self::Broadcast(shape) => {
                let x = front(x, axes[0])?;
                let missing = shape.len() - (x.ndim() - 1);
                let aligned: Vec<isize> = std::iter::once(b)
                    .chain(std::iter::repeat_n(1, missing))
                    .chain(x.dims()[1..].iter().copied())
                    .map(|n| n as isize)
                    .collect();
                let out = broadcast_to(&reshape(&x, &aligned)?, &prepend(b, shape))?;
                Ok((vec![out], 0))
            }
            Self::Reduce { op, axes: reduced } => at_front(Self::Reduce {
                op: *op,
                axes: reduced.iter().map(|a| a + 1).collect(),
            }),
            // Keys `[..., 2]`: a batch of keys draws one sample per key.
            Self::RandomBits { .. } => {
                if axes[0] == Some(x.ndim() - 1) {
                    return Err(crate::Error::InvalidArgument(
                        "random_bits: the key pair axis cannot be batched".into(),
                    ));
                }
                at_front(self.clone())
            }
            Self::ArgReduce { op, axis } => at_front(Self::ArgReduce {
                op: *op,
                axis: axis + 1,
            }),
            Self::Sort { axis } => at_front(Self::Sort { axis: axis + 1 }),
            Self::ArgSort { axis } => at_front(Self::ArgSort { axis: axis + 1 }),
            Self::Scan {
                op,
                axis,
                reverse,
                inclusive,
            } => at_front(Self::Scan {
                op: *op,
                axis: axis + 1,
                reverse: *reverse,
                inclusive: *inclusive,
            }),
            // Every operand gets the batch axis first (unbatched ones broadcast).
            Self::GatherAxis { axis } => {
                let all = (0..inputs.len())
                    .map(|i| front_full(&inputs[i], axes[i], b))
                    .collect::<Result<Vec<_>>>()?;
                Ok((record(Self::GatherAxis { axis: axis + 1 }, all)?, 0))
            }
            Self::ScatterAxis { axis, reduce } => {
                let all = (0..inputs.len())
                    .map(|i| front_full(&inputs[i], axes[i], b))
                    .collect::<Result<Vec<_>>>()?;
                let batched = Self::ScatterAxis {
                    axis: axis + 1,
                    reduce: *reduce,
                };
                Ok((record(batched, all)?, 0))
            }
            Self::Concatenate { axis } => {
                let all = (0..inputs.len())
                    .map(|i| front_full(&inputs[i], axes[i], b))
                    .collect::<Result<Vec<_>>>()?;
                Ok((record(Self::Concatenate { axis: axis + 1 }, all)?, 0))
            }
            // `matmul` broadcasts the size-1 batch axis of an unbatched operand.
            Self::Matmul => Ok((
                vec![matmul(&front(x, axes[0])?, &front(&inputs[1], axes[1])?)?],
                0,
            )),
            // `addmm` broadcasts `c` and size-1 batch axes, as `matmul` does.
            Self::AddMM { .. } => {
                let (alpha, beta) = self.addmm_factors();
                let a = front(x, axes[0])?;
                let b = front(&inputs[1], axes[1])?;
                let c = front(&inputs[2], axes[2])?;
                Ok((vec![crate::ops::imp::addmm(&c, &a, &b, alpha, beta)?], 0))
            }
            Self::Fast(op) => op.vmap(inputs, axes),
            Self::MetalKernel(_) => Err(crate::Error::InvalidArgument("metal_kernel: no derivative or batching rule (wrap the kernel in `custom_function`)".into())),
            Self::Compiled(fused) => fused.vmap(inputs, axes),
            Self::Quantize { quant } => {
                let (wq, scales, biases) = quantize(&front(x, axes[0])?, *quant)?;
                let mut out = vec![wq, scales];
                out.extend(biases);
                Ok((out, 0))
            }
            // Packed weights, scales and biases share their batch dimensions.
            Self::Dequantize { quant, dtype } => {
                let parts = (0..inputs.len())
                    .map(|i| front_full(&inputs[i], axes[i], b))
                    .collect::<Result<Vec<_>>>()?;
                Ok((
                    vec![dequantize(&parts[0], &parts[1], parts.get(2), *quant, Some(*dtype))?],
                    0,
                ))
            }
            // MLX has no rule. With unbatched weights, the batch joins `x`'s leading
            // dimensions; otherwise every operand gets the batch axis first, and 2-D weights
            // get size-1 axes for `x`'s other batch dimensions.
            Self::QuantizedMatmul { quant, transpose } => {
                let x = front(x, axes[0])?;
                let mut weights: Vec<Array> = inputs[1..].to_vec();
                if axes[1..].iter().any(Option::is_some) {
                    for (i, a) in weights.iter_mut().enumerate() {
                        *a = front_full(a, axes[i + 1], b)?;
                        while a.ndim() < x.ndim() {
                            *a = expand_dims(a, 1)?;
                        }
                    }
                }
                let out = quantized_matmul(
                    &x,
                    &weights[0],
                    &weights[1],
                    weights.get(2),
                    *transpose,
                    *quant,
                )?;
                Ok((vec![out], 0))
            }
            Self::Custom(call) => {
                let (primals, axes) = (&inputs[..call.primals], &axes[..call.primals]);
                match &call.rules.vmap {
                    Some(rule) => rule(primals, axes),
                    None => Ok((
                        crate::transforms::batch(&*call.rules.body, primals, axes)?,
                        0,
                    )),
                }
            }
            #[cfg(test)]
            Self::AddMultiply => self.elementwise(inputs, axes),
        }
    }

    /// Element-wise primitives: in place when every input is batched along the same axis
    /// with equal shapes; otherwise batch axes go first and inputs are broadcast together.
    fn elementwise(&self, inputs: &[Array], axes: &[Option<usize>]) -> Result<(Vec<Array>, usize)> {
        if let Some(axis) = axes[0]
            && axes.iter().all(|&a| a == Some(axis))
            && inputs.iter().all(|a| a.dims() == inputs[0].dims())
        {
            return Ok((record(self.clone(), inputs.to_vec())?, axis));
        }
        let fronted = inputs
            .iter()
            .zip(axes)
            .map(|(a, &ax)| front(a, ax))
            .collect::<Result<Vec<_>>>()?;
        let mut shape = fronted[0].dims().to_vec();
        for a in &fronted[1..] {
            shape = view::broadcast_shapes(self.name(), &shape, a.dims())?.to_vec();
        }
        let operands = fronted
            .iter()
            .map(|a| broadcast_to(a, &shape))
            .collect::<Result<Vec<_>>>()?;
        Ok((record(self.clone(), operands)?, 0))
    }
}