numr 0.5.2

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! Indexing operations for CUDA runtime

mod advanced;
mod argmax;
mod gather_scatter;
mod helpers;
mod masked;

use crate::error::Result;
use crate::ops::{IndexingOps, ScatterReduceOp};
use crate::runtime::cuda::{CudaClient, CudaRuntime};
use crate::tensor::Tensor;

impl IndexingOps<CudaRuntime> for CudaClient {
    fn argmax(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        keepdim: bool,
    ) -> Result<Tensor<CudaRuntime>> {
        argmax::argmax(self, a, dim, keepdim)
    }

    fn argmin(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        keepdim: bool,
    ) -> Result<Tensor<CudaRuntime>> {
        argmax::argmin(self, a, dim, keepdim)
    }

    fn gather(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        index: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::gather(self, a, dim, index)
    }

    fn scatter(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        index: &Tensor<CudaRuntime>,
        src: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::scatter(self, a, dim, index, src)
    }

    fn index_select(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        index: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::index_select(self, a, dim, index)
    }

    fn index_put(
        &self,
        a: &Tensor<CudaRuntime>,
        dim: usize,
        index: &Tensor<CudaRuntime>,
        src: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::index_put(self, a, dim, index, src)
    }

    fn masked_select(
        &self,
        a: &Tensor<CudaRuntime>,
        mask: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        masked::masked_select(self, a, mask)
    }

    fn masked_fill(
        &self,
        a: &Tensor<CudaRuntime>,
        mask: &Tensor<CudaRuntime>,
        value: f64,
    ) -> Result<Tensor<CudaRuntime>> {
        masked::masked_fill(self, a, mask, value)
    }

    fn embedding_lookup(
        &self,
        embeddings: &Tensor<CudaRuntime>,
        indices: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        advanced::embedding_lookup(self, embeddings, indices)
    }

    fn scatter_reduce(
        &self,
        dst: &Tensor<CudaRuntime>,
        dim: usize,
        index: &Tensor<CudaRuntime>,
        src: &Tensor<CudaRuntime>,
        op: ScatterReduceOp,
        include_self: bool,
    ) -> Result<Tensor<CudaRuntime>> {
        advanced::scatter_reduce(self, dst, dim, index, src, op, include_self)
    }

    fn gather_nd(
        &self,
        input: &Tensor<CudaRuntime>,
        indices: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        advanced::gather_nd(self, input, indices)
    }

    fn bincount(
        &self,
        input: &Tensor<CudaRuntime>,
        weights: Option<&Tensor<CudaRuntime>>,
        minlength: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        advanced::bincount(self, input, weights, minlength)
    }

    fn gather_2d(
        &self,
        input: &Tensor<CudaRuntime>,
        rows: &Tensor<CudaRuntime>,
        cols: &Tensor<CudaRuntime>,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::gather_2d(self, input, rows, cols)
    }

    fn slice_assign(
        &self,
        dst: &Tensor<CudaRuntime>,
        src: &Tensor<CudaRuntime>,
        dim: usize,
        start: usize,
    ) -> Result<Tensor<CudaRuntime>> {
        gather_scatter::slice_assign(self, dst, src, dim, start)
    }
}