1use super::*;
2
3impl QTensorOps<Self> for Dispatch {
4 fn q_from_data(data: TensorData, device: &Device<Self>) -> QuantizedTensor<Self> {
5 let selection = Selection::factory(device);
6 macro_rules! call {
7 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
8 let result = <$B>::q_from_data(data, $device);
9 pack!($mode, $variant, Quantized, result)
10 }};
11 }
12 dispatch!(selection, call)
13 }
14
15 fn quantize( tensor: FloatTensor<Self>, scheme: &QuantScheme, qparams: QuantizationParametersPrimitive<Self>, ) -> QuantizedTensor<Self> {
16 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&qparams.scales)));
17 macro_rules! call {
18 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
19 let result = <$B>::quantize(unpack!($mode, $variant, Float, tensor, $context), scheme, QuantizationParametersPrimitive { scales: unpack!($mode, $variant, Float, qparams.scales, $context) });
20 pack!($mode, $variant, Quantized, result)
21 }};
22 }
23 dispatch!(selection, call)
24 }
25
26 fn quantize_dynamic(tensor: FloatTensor<Self>, scheme: &QuantScheme) -> QuantizedTensor<Self> {
27 let selection = Selection::tensors(core::iter::once(&tensor));
28 macro_rules! call {
29 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
30 let result = <$B>::quantize_dynamic(unpack!($mode, $variant, Float, tensor, $context), scheme);
31 pack!($mode, $variant, Quantized, result)
32 }};
33 }
34 dispatch!(selection, call)
35 }
36
37 fn quantize_dynamic_with_precision(tensor: FloatTensor<Self>, scheme: &QuantScheme, calibration_dtype: FloatDType) -> QuantizedTensor<Self> {
38 let selection = Selection::tensors(core::iter::once(&tensor));
39 macro_rules! call {
40 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
41 let result = <$B>::quantize_dynamic_with_precision(unpack!($mode, $variant, Float, tensor, $context), scheme, calibration_dtype);
42 pack!($mode, $variant, Quantized, result)
43 }};
44 }
45 dispatch!(selection, call)
46 }
47
48 fn dequantize(tensor: QuantizedTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
49 let selection = Selection::tensors(core::iter::once(&tensor));
50 macro_rules! call {
51 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
52 let result = <$B>::dequantize(unpack!($mode, $variant, Quantized, tensor, $context), dtype);
53 pack!($mode, $variant, Float, result)
54 }};
55 }
56 dispatch!(selection, call)
57 }
58
59 fn q_device(tensor: &QuantizedTensor<Self>) -> Device<Self> {
60 tensor.device()
61 }
62
63 fn q_to_device(tensor: QuantizedTensor<Self>, device: &Device<Self>) -> QuantizedTensor<Self> {
64 crate::transfer::q_to_device(tensor, device)
65 }
66
67 fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
68 let selection = Selection::tensors(core::iter::once(&tensor));
69 macro_rules! call {
70 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
71 let result = <$B>::q_reshape(unpack!($mode, $variant, Quantized, tensor, $context), shape);
72 pack!($mode, $variant, Quantized, result)
73 }};
74 }
75 dispatch!(selection, call)
76 }
77
78 fn q_into_data( tensor: QuantizedTensor<Self>, ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send {
79 let selection = Selection::tensors(core::iter::once(&tensor));
80 macro_rules! call {
81 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
82 let future = <$B>::q_into_data(unpack!($mode, $variant, Quantized, tensor, $context));
83 let future: core::pin::Pin<Box<dyn Future<Output = Result<TensorData, ExecutionError>> + Send>> = Box::pin(async move { future.await });
84 future
85 }};
86 }
87 dispatch!(selection, call)
88 }
89
90 fn q_detach(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
91 let selection = Selection::tensors(core::iter::once(&tensor));
92 macro_rules! call {
93 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
94 let result = <$B>::q_detach(unpack!($mode, $variant, Quantized, tensor, $context));
95 pack!($mode, $variant, Quantized, result)
96 }};
97 }
98 dispatch!(selection, call)
99 }
100
101 fn q_set_require_grad(tensor: QuantizedTensor<Self>, _require_grad: bool) -> QuantizedTensor<Self> {
102 let selection = Selection::tensors(core::iter::once(&tensor));
103 macro_rules! call {
104 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
105 let result = <$B>::q_set_require_grad(unpack!($mode, $variant, Quantized, tensor, $context), _require_grad);
106 pack!($mode, $variant, Quantized, result)
107 }};
108 }
109 dispatch!(selection, call)
110 }
111
112 fn q_is_require_grad(_tensor: &QuantizedTensor<Self>) -> bool {
113 let selection = Selection::tensors(core::iter::once(_tensor));
114 macro_rules! call {
115 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
116 let result = <$B>::q_is_require_grad(borrow_tensor!($mode, $variant, Quantized, _tensor));
117 result
118 }};
119 }
120 dispatch!(selection, call)
121 }
122
123 fn q_is_tracked(tensor: &QuantizedTensor<Self>) -> bool {
124 let selection = Selection::tensors(core::iter::once(tensor));
125 macro_rules! call {
126 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
127 let result = <$B>::q_is_tracked(borrow_tensor!($mode, $variant, Quantized, tensor));
128 result
129 }};
130 }
131 dispatch!(selection, call)
132 }
133
134 fn q_gradient_checkpointing_strategy(_tensor: &QuantizedTensor<Self>) -> Option<ruda_tensor::GradientCheckpointingStrategy> {
135 match _tensor.context() { AutodiffContext::Disabled => None, AutodiffContext::Enabled(strategy) => Some(strategy) }
136 }
137
138 fn q_autodiff_context(tensor: &QuantizedTensor<Self>) -> ruda_tensor::AutodiffContext {
139 tensor.context()
140 }
141
142 fn q_with_autodiff_context(tensor: QuantizedTensor<Self>, context: ruda_tensor::AutodiffContext) -> QuantizedTensor<Self> {
143 tensor.with_context(context)
144 }
145
146 fn q_with_gradient_checkpointing_strategy( _tensor: QuantizedTensor<Self>, _strategy: ruda_tensor::GradientCheckpointingStrategy, ) -> QuantizedTensor<Self> {
147 assert!(matches!(_tensor.context(), AutodiffContext::Enabled(_)), "Checkpointing requires autodiff");
148 _tensor.with_context(AutodiffContext::Enabled(_strategy))
149 }
150
151 fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
152 let selection = Selection::tensors(core::iter::once(&tensor));
153 macro_rules! call {
154 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
155 let result = <$B>::q_expand(unpack!($mode, $variant, Quantized, tensor, $context), shape);
156 pack!($mode, $variant, Quantized, result)
157 }};
158 }
159 dispatch!(selection, call)
160 }
161
162 fn q_transpose(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
163 let selection = Selection::tensors(core::iter::once(&tensor));
164 macro_rules! call {
165 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
166 let result = <$B>::q_transpose(unpack!($mode, $variant, Quantized, tensor, $context));
167 pack!($mode, $variant, Quantized, result)
168 }};
169 }
170 dispatch!(selection, call)
171 }
172
173 fn q_swap_dims(tensor: QuantizedTensor<Self>, dim1: usize, dim2: usize) -> QuantizedTensor<Self> {
174 let selection = Selection::tensors(core::iter::once(&tensor));
175 macro_rules! call {
176 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
177 let result = <$B>::q_swap_dims(unpack!($mode, $variant, Quantized, tensor, $context), dim1, dim2);
178 pack!($mode, $variant, Quantized, result)
179 }};
180 }
181 dispatch!(selection, call)
182 }
183
184 fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
185 let selection = Selection::tensors(core::iter::once(&tensor));
186 macro_rules! call {
187 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
188 let result = <$B>::q_permute(unpack!($mode, $variant, Quantized, tensor, $context), axes);
189 pack!($mode, $variant, Quantized, result)
190 }};
191 }
192 dispatch!(selection, call)
193 }
194
195 fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
196 let selection = Selection::tensors(core::iter::once(&tensor));
197 macro_rules! call {
198 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
199 let result = <$B>::q_flip(unpack!($mode, $variant, Quantized, tensor, $context), axes);
200 pack!($mode, $variant, Quantized, result)
201 }};
202 }
203 dispatch!(selection, call)
204 }
205
206 fn q_select( tensor: QuantizedTensor<Self>, dim: usize, indices: IntTensor<Self>, ) -> QuantizedTensor<Self> {
207 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
208 macro_rules! call {
209 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
210 let result = <$B>::q_select(unpack!($mode, $variant, Quantized, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context));
211 pack!($mode, $variant, Quantized, result)
212 }};
213 }
214 dispatch!(selection, call)
215 }
216
217 fn q_slice(tensor: QuantizedTensor<Self>, slices: &[Slice]) -> QuantizedTensor<Self> {
218 let selection = Selection::tensors(core::iter::once(&tensor));
219 macro_rules! call {
220 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
221 let result = <$B>::q_slice(unpack!($mode, $variant, Quantized, tensor, $context), slices);
222 pack!($mode, $variant, Quantized, result)
223 }};
224 }
225 dispatch!(selection, call)
226 }
227
228 fn q_gather( dim: usize, tensor: QuantizedTensor<Self>, indices: IntTensor<Self>, ) -> QuantizedTensor<Self> {
229 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
230 macro_rules! call {
231 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
232 let result = <$B>::q_gather(dim, unpack!($mode, $variant, Quantized, tensor, $context), unpack!($mode, $variant, Int, indices, $context));
233 pack!($mode, $variant, Quantized, result)
234 }};
235 }
236 dispatch!(selection, call)
237 }
238
239 fn q_repeat_dim(tensor: QuantizedTensor<Self>, dim: usize, times: usize) -> QuantizedTensor<Self> {
240 let selection = Selection::tensors(core::iter::once(&tensor));
241 macro_rules! call {
242 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
243 let result = <$B>::q_repeat_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim, times);
244 pack!($mode, $variant, Quantized, result)
245 }};
246 }
247 dispatch!(selection, call)
248 }
249
250 fn q_add(lhs: QuantizedTensor<Self>, rhs: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
251 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
252 macro_rules! call {
253 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
254 let result = <$B>::q_add(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Quantized, rhs, $context));
255 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
256 }};
257 }
258 dispatch!(selection, call)
259 }
260
261 fn q_add_scalar(lhs: QuantizedTensor<Self>, rhs: Scalar) -> TensorPrimitive<Self> {
262 let selection = Selection::tensors(core::iter::once(&lhs));
263 macro_rules! call {
264 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
265 let result = <$B>::q_add_scalar(unpack!($mode, $variant, Quantized, lhs, $context), rhs);
266 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
267 }};
268 }
269 dispatch!(selection, call)
270 }
271
272 fn q_clamp_min(tensor: QuantizedTensor<Self>, min: Scalar) -> TensorPrimitive<Self> {
273 let selection = Selection::tensors(core::iter::once(&tensor));
274 macro_rules! call {
275 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
276 let result = <$B>::q_clamp_min(unpack!($mode, $variant, Quantized, tensor, $context), min);
277 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
278 }};
279 }
280 dispatch!(selection, call)
281 }
282
283 fn q_clamp_max(tensor: QuantizedTensor<Self>, max: Scalar) -> TensorPrimitive<Self> {
284 let selection = Selection::tensors(core::iter::once(&tensor));
285 macro_rules! call {
286 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
287 let result = <$B>::q_clamp_max(unpack!($mode, $variant, Quantized, tensor, $context), max);
288 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
289 }};
290 }
291 dispatch!(selection, call)
292 }
293
294 fn q_clamp(tensor: QuantizedTensor<Self>, min: Scalar, max: Scalar) -> TensorPrimitive<Self> {
295 let selection = Selection::tensors(core::iter::once(&tensor));
296 macro_rules! call {
297 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
298 let result = <$B>::q_clamp(unpack!($mode, $variant, Quantized, tensor, $context), min, max);
299 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
300 }};
301 }
302 dispatch!(selection, call)
303 }
304
305 fn q_sub(lhs: QuantizedTensor<Self>, rhs: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
306 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
307 macro_rules! call {
308 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
309 let result = <$B>::q_sub(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Quantized, rhs, $context));
310 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
311 }};
312 }
313 dispatch!(selection, call)
314 }
315
316 fn q_sub_scalar(lhs: QuantizedTensor<Self>, rhs: Scalar) -> TensorPrimitive<Self> {
317 let selection = Selection::tensors(core::iter::once(&lhs));
318 macro_rules! call {
319 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
320 let result = <$B>::q_sub_scalar(unpack!($mode, $variant, Quantized, lhs, $context), rhs);
321 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
322 }};
323 }
324 dispatch!(selection, call)
325 }
326
327 fn q_mul(lhs: QuantizedTensor<Self>, rhs: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
328 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
329 macro_rules! call {
330 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
331 let result = <$B>::q_mul(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Quantized, rhs, $context));
332 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
333 }};
334 }
335 dispatch!(selection, call)
336 }
337
338 fn q_mul_scalar(lhs: QuantizedTensor<Self>, rhs: Scalar) -> TensorPrimitive<Self> {
339 let selection = Selection::tensors(core::iter::once(&lhs));
340 macro_rules! call {
341 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
342 let result = <$B>::q_mul_scalar(unpack!($mode, $variant, Quantized, lhs, $context), rhs);
343 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
344 }};
345 }
346 dispatch!(selection, call)
347 }
348
349 fn q_div(lhs: QuantizedTensor<Self>, rhs: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
350 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
351 macro_rules! call {
352 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
353 let result = <$B>::q_div(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Quantized, rhs, $context));
354 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
355 }};
356 }
357 dispatch!(selection, call)
358 }
359
360 fn q_div_scalar(lhs: QuantizedTensor<Self>, rhs: Scalar) -> TensorPrimitive<Self> {
361 let selection = Selection::tensors(core::iter::once(&lhs));
362 macro_rules! call {
363 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
364 let result = <$B>::q_div_scalar(unpack!($mode, $variant, Quantized, lhs, $context), rhs);
365 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
366 }};
367 }
368 dispatch!(selection, call)
369 }
370
371 fn q_matmul(lhs: TensorPrimitive<Self>, rhs: TensorPrimitive<Self>) -> TensorPrimitive<Self> {
372 let selection = Selection::tensors(core::iter::once(match &lhs { TensorPrimitive::Float(tensor) | TensorPrimitive::QFloat(tensor) => tensor }).chain(core::iter::once(match &rhs { TensorPrimitive::Float(tensor) | TensorPrimitive::QFloat(tensor) => tensor })));
373 macro_rules! call {
374 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
375 let result = <$B>::q_matmul(match lhs { TensorPrimitive::Float(value) => TensorPrimitive::Float(unpack!($mode, $variant, Float, value, $context)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(unpack!($mode, $variant, Quantized, value, $context)) }, match rhs { TensorPrimitive::Float(value) => TensorPrimitive::Float(unpack!($mode, $variant, Float, value, $context)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(unpack!($mode, $variant, Quantized, value, $context)) });
376 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
377 }};
378 }
379 dispatch!(selection, call)
380 }
381
382 fn q_matmul_default(lhs: TensorPrimitive<Self>, rhs: TensorPrimitive<Self>) -> TensorPrimitive<Self> {
383 let selection = Selection::tensors(core::iter::once(match &lhs { TensorPrimitive::Float(tensor) | TensorPrimitive::QFloat(tensor) => tensor }).chain(core::iter::once(match &rhs { TensorPrimitive::Float(tensor) | TensorPrimitive::QFloat(tensor) => tensor })));
384 macro_rules! call {
385 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
386 let result = <$B>::q_matmul_default(match lhs { TensorPrimitive::Float(value) => TensorPrimitive::Float(unpack!($mode, $variant, Float, value, $context)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(unpack!($mode, $variant, Quantized, value, $context)) }, match rhs { TensorPrimitive::Float(value) => TensorPrimitive::Float(unpack!($mode, $variant, Float, value, $context)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(unpack!($mode, $variant, Quantized, value, $context)) });
387 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
388 }};
389 }
390 dispatch!(selection, call)
391 }
392
393 fn q_neg(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
394 let selection = Selection::tensors(core::iter::once(&tensor));
395 macro_rules! call {
396 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
397 let result = <$B>::q_neg(unpack!($mode, $variant, Quantized, tensor, $context));
398 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
399 }};
400 }
401 dispatch!(selection, call)
402 }
403
404 fn q_recip(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
405 let selection = Selection::tensors(core::iter::once(&tensor));
406 macro_rules! call {
407 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
408 let result = <$B>::q_recip(unpack!($mode, $variant, Quantized, tensor, $context));
409 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
410 }};
411 }
412 dispatch!(selection, call)
413 }
414
415 fn q_sum(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
416 let selection = Selection::tensors(core::iter::once(&tensor));
417 macro_rules! call {
418 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
419 let result = <$B>::q_sum(unpack!($mode, $variant, Quantized, tensor, $context));
420 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
421 }};
422 }
423 dispatch!(selection, call)
424 }
425
426 fn q_sum_dim(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
427 let selection = Selection::tensors(core::iter::once(&tensor));
428 macro_rules! call {
429 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
430 let result = <$B>::q_sum_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
431 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
432 }};
433 }
434 dispatch!(selection, call)
435 }
436
437 fn q_prod(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
438 let selection = Selection::tensors(core::iter::once(&tensor));
439 macro_rules! call {
440 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
441 let result = <$B>::q_prod(unpack!($mode, $variant, Quantized, tensor, $context));
442 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
443 }};
444 }
445 dispatch!(selection, call)
446 }
447
448 fn q_prod_dim(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
449 let selection = Selection::tensors(core::iter::once(&tensor));
450 macro_rules! call {
451 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
452 let result = <$B>::q_prod_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
453 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
454 }};
455 }
456 dispatch!(selection, call)
457 }
458
459 fn q_mean(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
460 let selection = Selection::tensors(core::iter::once(&tensor));
461 macro_rules! call {
462 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
463 let result = <$B>::q_mean(unpack!($mode, $variant, Quantized, tensor, $context));
464 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
465 }};
466 }
467 dispatch!(selection, call)
468 }
469
470 fn q_mean_dim(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
471 let selection = Selection::tensors(core::iter::once(&tensor));
472 macro_rules! call {
473 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
474 let result = <$B>::q_mean_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
475 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
476 }};
477 }
478 dispatch!(selection, call)
479 }
480
481 fn q_cumsum(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
482 let selection = Selection::tensors(core::iter::once(&tensor));
483 macro_rules! call {
484 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
485 let result = <$B>::q_cumsum(unpack!($mode, $variant, Quantized, tensor, $context), dim);
486 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
487 }};
488 }
489 dispatch!(selection, call)
490 }
491
492 fn q_cumprod(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
493 let selection = Selection::tensors(core::iter::once(&tensor));
494 macro_rules! call {
495 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
496 let result = <$B>::q_cumprod(unpack!($mode, $variant, Quantized, tensor, $context), dim);
497 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
498 }};
499 }
500 dispatch!(selection, call)
501 }
502
503 fn q_cummin(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
504 let selection = Selection::tensors(core::iter::once(&tensor));
505 macro_rules! call {
506 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
507 let result = <$B>::q_cummin(unpack!($mode, $variant, Quantized, tensor, $context), dim);
508 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
509 }};
510 }
511 dispatch!(selection, call)
512 }
513
514 fn q_cummax(tensor: QuantizedTensor<Self>, dim: usize) -> TensorPrimitive<Self> {
515 let selection = Selection::tensors(core::iter::once(&tensor));
516 macro_rules! call {
517 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
518 let result = <$B>::q_cummax(unpack!($mode, $variant, Quantized, tensor, $context), dim);
519 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
520 }};
521 }
522 dispatch!(selection, call)
523 }
524
525 fn q_exp(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
526 let selection = Selection::tensors(core::iter::once(&tensor));
527 macro_rules! call {
528 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
529 let result = <$B>::q_exp(unpack!($mode, $variant, Quantized, tensor, $context));
530 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
531 }};
532 }
533 dispatch!(selection, call)
534 }
535
536 fn q_log(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
537 let selection = Selection::tensors(core::iter::once(&tensor));
538 macro_rules! call {
539 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
540 let result = <$B>::q_log(unpack!($mode, $variant, Quantized, tensor, $context));
541 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
542 }};
543 }
544 dispatch!(selection, call)
545 }
546
547 fn q_log1p(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
548 let selection = Selection::tensors(core::iter::once(&tensor));
549 macro_rules! call {
550 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
551 let result = <$B>::q_log1p(unpack!($mode, $variant, Quantized, tensor, $context));
552 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
553 }};
554 }
555 dispatch!(selection, call)
556 }
557
558 fn q_powf(lhs: QuantizedTensor<Self>, rhs: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
559 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
560 macro_rules! call {
561 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
562 let result = <$B>::q_powf(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Quantized, rhs, $context));
563 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
564 }};
565 }
566 dispatch!(selection, call)
567 }
568
569 fn q_powi(lhs: QuantizedTensor<Self>, rhs: IntTensor<Self>) -> TensorPrimitive<Self> {
570 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
571 macro_rules! call {
572 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
573 let result = <$B>::q_powi(unpack!($mode, $variant, Quantized, lhs, $context), unpack!($mode, $variant, Int, rhs, $context));
574 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
575 }};
576 }
577 dispatch!(selection, call)
578 }
579
580 fn q_powi_scalar(lhs: QuantizedTensor<Self>, rhs: Scalar) -> TensorPrimitive<Self> {
581 let selection = Selection::tensors(core::iter::once(&lhs));
582 macro_rules! call {
583 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
584 let result = <$B>::q_powi_scalar(unpack!($mode, $variant, Quantized, lhs, $context), rhs);
585 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
586 }};
587 }
588 dispatch!(selection, call)
589 }
590
591 fn q_powf_scalar(tensor: QuantizedTensor<Self>, value: Scalar) -> TensorPrimitive<Self> {
592 let selection = Selection::tensors(core::iter::once(&tensor));
593 macro_rules! call {
594 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
595 let result = <$B>::q_powf_scalar(unpack!($mode, $variant, Quantized, tensor, $context), value);
596 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
597 }};
598 }
599 dispatch!(selection, call)
600 }
601
602 fn q_sqrt(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
603 let selection = Selection::tensors(core::iter::once(&tensor));
604 macro_rules! call {
605 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
606 let result = <$B>::q_sqrt(unpack!($mode, $variant, Quantized, tensor, $context));
607 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
608 }};
609 }
610 dispatch!(selection, call)
611 }
612
613 fn q_abs(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
614 let selection = Selection::tensors(core::iter::once(&tensor));
615 macro_rules! call {
616 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
617 let result = <$B>::q_abs(unpack!($mode, $variant, Quantized, tensor, $context));
618 pack!($mode, $variant, Quantized, result)
619 }};
620 }
621 dispatch!(selection, call)
622 }
623
624 fn q_cos(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
625 let selection = Selection::tensors(core::iter::once(&tensor));
626 macro_rules! call {
627 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
628 let result = <$B>::q_cos(unpack!($mode, $variant, Quantized, tensor, $context));
629 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
630 }};
631 }
632 dispatch!(selection, call)
633 }
634
635 fn q_sin(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
636 let selection = Selection::tensors(core::iter::once(&tensor));
637 macro_rules! call {
638 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
639 let result = <$B>::q_sin(unpack!($mode, $variant, Quantized, tensor, $context));
640 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
641 }};
642 }
643 dispatch!(selection, call)
644 }
645
646 fn q_tan(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
647 let selection = Selection::tensors(core::iter::once(&tensor));
648 macro_rules! call {
649 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
650 let result = <$B>::q_tan(unpack!($mode, $variant, Quantized, tensor, $context));
651 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
652 }};
653 }
654 dispatch!(selection, call)
655 }
656
657 fn q_cosh(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
658 let selection = Selection::tensors(core::iter::once(&tensor));
659 macro_rules! call {
660 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
661 let result = <$B>::q_cosh(unpack!($mode, $variant, Quantized, tensor, $context));
662 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
663 }};
664 }
665 dispatch!(selection, call)
666 }
667
668 fn q_sinh(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
669 let selection = Selection::tensors(core::iter::once(&tensor));
670 macro_rules! call {
671 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
672 let result = <$B>::q_sinh(unpack!($mode, $variant, Quantized, tensor, $context));
673 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
674 }};
675 }
676 dispatch!(selection, call)
677 }
678
679 fn q_tanh(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
680 let selection = Selection::tensors(core::iter::once(&tensor));
681 macro_rules! call {
682 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
683 let result = <$B>::q_tanh(unpack!($mode, $variant, Quantized, tensor, $context));
684 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
685 }};
686 }
687 dispatch!(selection, call)
688 }
689
690 fn q_erf(tensor: QuantizedTensor<Self>) -> TensorPrimitive<Self> {
691 let selection = Selection::tensors(core::iter::once(&tensor));
692 macro_rules! call {
693 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
694 let result = <$B>::q_erf(unpack!($mode, $variant, Quantized, tensor, $context));
695 match result { TensorPrimitive::Float(value) => TensorPrimitive::Float(pack!($mode, $variant, Float, value)), TensorPrimitive::QFloat(value) => TensorPrimitive::QFloat(pack!($mode, $variant, Quantized, value)) }
696 }};
697 }
698 dispatch!(selection, call)
699 }
700
701 fn q_cat(tensors: Vec<QuantizedTensor<Self>>, dim: usize) -> QuantizedTensor<Self> {
702 let selection = Selection::tensors(tensors.iter());
703 macro_rules! call {
704 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
705 let result = <$B>::q_cat(tensors.into_iter().map(|value| unpack!($mode, $variant, Quantized, value, $context)).collect(), dim);
706 pack!($mode, $variant, Quantized, result)
707 }};
708 }
709 dispatch!(selection, call)
710 }
711
712 fn q_argmax(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
713 let selection = Selection::tensors(core::iter::once(&tensor));
714 macro_rules! call {
715 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
716 let result = <$B>::q_argmax(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
717 pack!($mode, $variant, Int, result)
718 }};
719 }
720 dispatch!(selection, call)
721 }
722
723 fn q_argtopk( tensor: QuantizedTensor<Self>, dim: usize, k: usize, out_dtype: IntDType, ) -> IntTensor<Self> {
724 let selection = Selection::tensors(core::iter::once(&tensor));
725 macro_rules! call {
726 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
727 let result = <$B>::q_argtopk(unpack!($mode, $variant, Quantized, tensor, $context), dim, k, out_dtype);
728 pack!($mode, $variant, Int, result)
729 }};
730 }
731 dispatch!(selection, call)
732 }
733
734 fn q_topk(tensor: QuantizedTensor<Self>, dim: usize, k: usize) -> QuantizedTensor<Self> {
735 let selection = Selection::tensors(core::iter::once(&tensor));
736 macro_rules! call {
737 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
738 let result = <$B>::q_topk(unpack!($mode, $variant, Quantized, tensor, $context), dim, k);
739 pack!($mode, $variant, Quantized, result)
740 }};
741 }
742 dispatch!(selection, call)
743 }
744
745 fn q_argmin(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
746 let selection = Selection::tensors(core::iter::once(&tensor));
747 macro_rules! call {
748 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
749 let result = <$B>::q_argmin(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
750 pack!($mode, $variant, Int, result)
751 }};
752 }
753 dispatch!(selection, call)
754 }
755
756 fn q_max(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
757 let selection = Selection::tensors(core::iter::once(&tensor));
758 macro_rules! call {
759 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
760 let result = <$B>::q_max(unpack!($mode, $variant, Quantized, tensor, $context));
761 pack!($mode, $variant, Quantized, result)
762 }};
763 }
764 dispatch!(selection, call)
765 }
766
767 fn q_max_dim(tensor: QuantizedTensor<Self>, dim: usize) -> QuantizedTensor<Self> {
768 let selection = Selection::tensors(core::iter::once(&tensor));
769 macro_rules! call {
770 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
771 let result = <$B>::q_max_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
772 pack!($mode, $variant, Quantized, result)
773 }};
774 }
775 dispatch!(selection, call)
776 }
777
778 fn q_max_dim_with_indices( tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType, ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
779 let selection = Selection::tensors(core::iter::once(&tensor));
780 macro_rules! call {
781 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
782 let result = <$B>::q_max_dim_with_indices(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
783 { let (value0, value1) = result; (pack!($mode, $variant, Quantized, value0), pack!($mode, $variant, Int, value1)) }
784 }};
785 }
786 dispatch!(selection, call)
787 }
788
789 fn q_min(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
790 let selection = Selection::tensors(core::iter::once(&tensor));
791 macro_rules! call {
792 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
793 let result = <$B>::q_min(unpack!($mode, $variant, Quantized, tensor, $context));
794 pack!($mode, $variant, Quantized, result)
795 }};
796 }
797 dispatch!(selection, call)
798 }
799
800 fn q_min_dim(tensor: QuantizedTensor<Self>, dim: usize) -> QuantizedTensor<Self> {
801 let selection = Selection::tensors(core::iter::once(&tensor));
802 macro_rules! call {
803 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
804 let result = <$B>::q_min_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
805 pack!($mode, $variant, Quantized, result)
806 }};
807 }
808 dispatch!(selection, call)
809 }
810
811 fn q_min_dim_with_indices( tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType, ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
812 let selection = Selection::tensors(core::iter::once(&tensor));
813 macro_rules! call {
814 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
815 let result = <$B>::q_min_dim_with_indices(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
816 { let (value0, value1) = result; (pack!($mode, $variant, Quantized, value0), pack!($mode, $variant, Int, value1)) }
817 }};
818 }
819 dispatch!(selection, call)
820 }
821
822 fn q_max_abs(tensor: QuantizedTensor<Self>) -> QuantizedTensor<Self> {
823 let selection = Selection::tensors(core::iter::once(&tensor));
824 macro_rules! call {
825 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
826 let result = <$B>::q_max_abs(unpack!($mode, $variant, Quantized, tensor, $context));
827 pack!($mode, $variant, Quantized, result)
828 }};
829 }
830 dispatch!(selection, call)
831 }
832
833 fn q_max_abs_dim(tensor: QuantizedTensor<Self>, dim: usize) -> QuantizedTensor<Self> {
834 let selection = Selection::tensors(core::iter::once(&tensor));
835 macro_rules! call {
836 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
837 let result = <$B>::q_max_abs_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim);
838 pack!($mode, $variant, Quantized, result)
839 }};
840 }
841 dispatch!(selection, call)
842 }
843
844 fn q_any(tensor: QuantizedTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
845 let selection = Selection::tensors(core::iter::once(&tensor));
846 macro_rules! call {
847 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
848 let result = <$B>::q_any(unpack!($mode, $variant, Quantized, tensor, $context), out_dtype);
849 pack!($mode, $variant, Bool, result)
850 }};
851 }
852 dispatch!(selection, call)
853 }
854
855 fn q_any_dim(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
856 let selection = Selection::tensors(core::iter::once(&tensor));
857 macro_rules! call {
858 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
859 let result = <$B>::q_any_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
860 pack!($mode, $variant, Bool, result)
861 }};
862 }
863 dispatch!(selection, call)
864 }
865
866 fn q_all(tensor: QuantizedTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
867 let selection = Selection::tensors(core::iter::once(&tensor));
868 macro_rules! call {
869 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
870 let result = <$B>::q_all(unpack!($mode, $variant, Quantized, tensor, $context), out_dtype);
871 pack!($mode, $variant, Bool, result)
872 }};
873 }
874 dispatch!(selection, call)
875 }
876
877 fn q_all_dim(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
878 let selection = Selection::tensors(core::iter::once(&tensor));
879 macro_rules! call {
880 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
881 let result = <$B>::q_all_dim(unpack!($mode, $variant, Quantized, tensor, $context), dim, out_dtype);
882 pack!($mode, $variant, Bool, result)
883 }};
884 }
885 dispatch!(selection, call)
886 }
887
888 fn q_sort(tensor: QuantizedTensor<Self>, dim: usize, descending: bool) -> QuantizedTensor<Self> {
889 let selection = Selection::tensors(core::iter::once(&tensor));
890 macro_rules! call {
891 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
892 let result = <$B>::q_sort(unpack!($mode, $variant, Quantized, tensor, $context), dim, descending);
893 pack!($mode, $variant, Quantized, result)
894 }};
895 }
896 dispatch!(selection, call)
897 }
898
899 fn q_sort_with_indices( tensor: QuantizedTensor<Self>, dim: usize, descending: bool, out_dtype: IntDType, ) -> (QuantizedTensor<Self>, IntTensor<Self>) {
900 let selection = Selection::tensors(core::iter::once(&tensor));
901 macro_rules! call {
902 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
903 let result = <$B>::q_sort_with_indices(unpack!($mode, $variant, Quantized, tensor, $context), dim, descending, out_dtype);
904 { let (value0, value1) = result; (pack!($mode, $variant, Quantized, value0), pack!($mode, $variant, Int, value1)) }
905 }};
906 }
907 dispatch!(selection, call)
908 }
909
910 fn q_argsort( tensor: QuantizedTensor<Self>, dim: usize, descending: bool, out_dtype: IntDType, ) -> IntTensor<Self> {
911 let selection = Selection::tensors(core::iter::once(&tensor));
912 macro_rules! call {
913 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
914 let result = <$B>::q_argsort(unpack!($mode, $variant, Quantized, tensor, $context), dim, descending, out_dtype);
915 pack!($mode, $variant, Int, result)
916 }};
917 }
918 dispatch!(selection, call)
919 }
920}