tenferro-gpu 0.3.0

CubeCL-backed CUDA and WebGPU provider backends for tenferro tensors.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use crate::kernels::{CubeclKernelError, Result};

/// Validate that a reduction axis is in bounds for a tensor rank.
pub(crate) fn validate_axis(rank: usize, axis: usize) -> Result<()> {
    if axis >= rank {
        return Err(CubeclKernelError::InvalidAxis { axis, rank });
    }
    Ok(())
}

/// Return the keepdims output shape for reducing a single axis.
pub(crate) fn keepdims_output_shape(input_shape: &[usize], axis: usize) -> Result<Vec<usize>> {
    validate_axis(input_shape.len(), axis)?;
    let mut output = input_shape.to_vec();
    output[axis] = 1;
    Ok(output)
}

/// Validate that an output shape matches a single-axis keepdims reduction.
pub(crate) fn validate_keepdims_output_shape(
    input_shape: &[usize],
    output_shape: &[usize],
    axis: usize,
) -> Result<()> {
    let expected = keepdims_output_shape(input_shape, axis)?;
    if output_shape != expected {
        return Err(CubeclKernelError::MismatchOutputShape {
            expected,
            actual: output_shape.to_vec(),
        });
    }
    Ok(())
}