ruPRIM-host 0.1.19

ruPRIM CPU tensor primitives, reductions, indexing and SIMD kernels.
Documentation 
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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

use ruda_core::{bytes::Bytes, tensor::{DType, Shape, host::{HostTensor, Layout}}};
use crate::unary;

pub fn float_max_dim_with_indices(
    tensor: HostTensor,
    dim: usize,
    indices_dtype: ruda_core::tensor::IntDType,
) -> (HostTensor, HostTensor) {
    let (values, indices) = crate::reduce::max_dim_with_indices(tensor, dim);
    if indices.dtype() != DType::from(indices_dtype) {
        (values, crate::cast::int_cast(indices, indices_dtype))
    } else {
        (values, indices)
    }
}

pub fn float_min_dim_with_indices(
    tensor: HostTensor,
    dim: usize,
    indices_dtype: ruda_core::tensor::IntDType,
) -> (HostTensor, HostTensor) {
    let (values, indices) = crate::reduce::min_dim_with_indices(tensor, dim);
    if indices.dtype() != DType::from(indices_dtype) {
        (values, crate::cast::int_cast(indices, indices_dtype))
    } else {
        (values, indices)
    }
}

pub fn float_argmax(
    tensor: HostTensor,
    dim: usize,
    out_dtype: ruda_core::tensor::IntDType,
) -> HostTensor {
    let result = crate::reduce::argmax(tensor, dim);
    if result.dtype() != DType::from(out_dtype) {
        crate::cast::int_cast(result, out_dtype)
    } else {
        result
    }
}

pub fn float_argmin(
    tensor: HostTensor,
    dim: usize,
    out_dtype: ruda_core::tensor::IntDType,
) -> HostTensor {
    let result = crate::reduce::argmin(tensor, dim);
    if result.dtype() != DType::from(out_dtype) {
        crate::cast::int_cast(result, out_dtype)
    } else {
        result
    }
}

pub fn float_max_abs(tensor: HostTensor) -> HostTensor {
    let abs = unary::abs(tensor);
    crate::reduce::max(abs)
}

pub fn float_max_abs_dim(tensor: HostTensor, dim: usize) -> HostTensor {
    let abs = unary::abs(tensor);
    crate::reduce::max_dim(abs, dim)
}

pub fn int_mean(tensor: HostTensor) -> HostTensor {
    let n = tensor.layout().num_elements();
    assert!(n > 0, "int_mean: cannot take mean of empty tensor");
    let dtype = tensor.dtype();
    let sum_result = crate::reduce::sum(tensor);
    // Compute in i64 to avoid truncation of n for small int types
    macro_rules! compute_mean {
        ($ty:ty) => {{
            let data: &[$ty] = sum_result.storage();
            let mean_val = (data[0] as i64 / n as i64) as $ty;
            HostTensor::new(
                Bytes::from_elems(alloc::vec![mean_val]),
                Layout::contiguous(Shape::from(alloc::vec![1])),
                dtype,
            )
        }};
    }
    match dtype {
        DType::I64 => compute_mean!(i64),
        DType::I32 => compute_mean!(i32),
        DType::I16 => compute_mean!(i16),
        DType::I8 => compute_mean!(i8),
        other => panic!("int_mean: unsupported dtype {:?}", other),
    }
}

pub fn int_max_abs(tensor: HostTensor) -> HostTensor {
    let abs = crate::unary::int_abs(tensor);
    crate::reduce::max(abs)
}

pub fn int_max_abs_dim(tensor: HostTensor, dim: usize) -> HostTensor {
    let abs = crate::unary::int_abs(tensor);
    crate::reduce::max_dim(abs, dim)
}