forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
TensorCompare.cpp
38 lines (32 loc) · 958 Bytes
/
TensorCompare.cpp
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
#include <ATen/ATen.h>
#include <ATen/CPUApplyUtils.h>
#include <ATen/Dispatch.h>
#include <ATen/ExpandUtils.h>
#include <ATen/NativeFunctions.h>
#include <ATen/native/ReduceOpsUtils.h>
#include <ATen/native/TensorCompare.h>
#include <c10/util/Exception.h>
namespace at {
namespace native {
Tensor max_quant(const Tensor& self) {
return std::get<0>(self.reshape({-1}).max(/*dim=*/0));
}
Tensor min_quant(const Tensor& self) {
return std::get<0>(self.reshape({-1}).min(/*dim=*/0));
}
// TODO: move to TensorMath.cpp
std::tuple<Tensor, Tensor> sort_quant(
const Tensor& self,
int64_t dim,
bool descending) {
Tensor sort_int;
Tensor sort_indicies;
std::tie(sort_int, sort_indicies) =
at::sort(self.int_repr(), dim, descending);
return std::forward_as_tuple(
at::_make_per_tensor_quantized_tensor(
sort_int, self.q_scale(), self.q_zero_point()),
sort_indicies);
}
} // namespace native
} // namespace at