hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
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

/// argmax kernel
macro_rules! argmax_kernel {
    (
        $_:expr,
        $iterator:ident,
        $loop_size:ident,
        $loop_size2:ident,
        $result_ptr:ident,
        $a_ptr:ident,
        $a_last_stride:ident,
        $shape_len:ident
    ) => {
        let mut max_val = T::NEG_INF;
        let mut max_index = 0;
        for i in 0..$loop_size {
            let a_val = $a_ptr[i * $a_last_stride];
            if a_val._gt(max_val) {
                max_val = a_val;
                max_index = i;
            }
        }
        for j in (0..$shape_len - 1).rev() {
            if $iterator.prg[j as usize] < $iterator.a_shape[j as usize] {
                $iterator.prg[j as usize] += 1;
                $a_ptr.offset($iterator.strides[j as usize]);
                break;
            } else {
                $iterator.prg[j as usize] = 0;
                $a_ptr.offset(-$iterator.strides[j as usize] * $iterator.a_shape[j as usize]);
            }
        }
        unsafe { $result_ptr.ptr.write(max_index.cast()) };
        $result_ptr.add(1);
    };
}

/// argmin kernel
macro_rules! argmin_kernel {
    (
        $_:expr,
        $iterator:ident,
        $loop_size:ident,
        $loop_size2:ident,
        $result_ptr:ident,
        $a_ptr:ident,
        $a_last_stride:ident,
        $shape_len:ident
    ) => {
        let mut max_val = T::INF;
        let mut max_index = 0;
        for i in 0..$loop_size {
            let a_val = $a_ptr[i * $a_last_stride];
            if a_val._lt(max_val) {
                max_val = a_val;
                max_index = i;
            }
        }
        for j in (0..$shape_len - 1).rev() {
            if $iterator.prg[j as usize] < $iterator.a_shape[j as usize] {
                $iterator.prg[j as usize] += 1;
                $a_ptr.offset($iterator.strides[j as usize]);
                break;
            } else {
                $iterator.prg[j as usize] = 0;
                $a_ptr.offset(-$iterator.strides[j as usize] * $iterator.a_shape[j as usize]);
            }
        }
        unsafe { $result_ptr.ptr.write(max_index.cast()) };
        $result_ptr.add(1);
    };
}

pub(crate) use argmax_kernel;
pub(crate) use argmin_kernel;