ruPRIM-host 0.1.27

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

use ruda_core::tensor::{DType, host::HostTensor};

pub fn float_argtopk(
    tensor: HostTensor,
    dim: usize,
    k: usize,
    out_dtype: ruda_core::tensor::IntDType,
) -> HostTensor {
    let indices = super::argtopk(tensor, dim, k);
    if indices.dtype() != DType::from(out_dtype) {
        crate::cast::int_cast(indices, out_dtype)
    } else {
        indices
    }
}

pub fn int_argtopk(tensor: HostTensor, dim: usize, k: usize) -> HostTensor {
    let dtype = tensor.dtype();
    let indices = super::argtopk(tensor, dim, k);
    if indices.dtype() != dtype {
        crate::cast::int_cast(indices, dtype.into())
    } else {
        indices
    }
}

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

pub fn float_argsort(
    tensor: HostTensor,
    dim: usize,
    descending: bool,
    out_dtype: ruda_core::tensor::IntDType,
) -> HostTensor {
    let indices = crate::sort::argsort(tensor, dim, descending);
    if indices.dtype() != DType::from(out_dtype) {
        crate::cast::int_cast(indices, out_dtype)
    } else {
        indices
    }
}