1 // Returns the frequency of elements of input non-negative integer tensor.
2 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
3
4 #include <ATen/core/Tensor.h>
5 #include <ATen/Dispatch.h>
6 #include <c10/util/irange.h>
7
8 #ifndef AT_PER_OPERATOR_HEADERS
9 #include <ATen/Functions.h>
10 #include <ATen/NativeFunctions.h>
11 #else
12 #include <ATen/ops/bincount_native.h>
13 #include <ATen/ops/zeros.h>
14 #endif
15
16 namespace at::native {
17
18 ///////////////// bincount /////////////////
19 namespace {
20
21 template <typename input_t, typename weights_t>
_bincount_cpu_template(const Tensor & self,const Tensor & weights,int64_t minlength)22 Tensor _bincount_cpu_template(
23 const Tensor& self,
24 const Tensor& weights,
25 int64_t minlength) {
26 if (minlength < 0) {
27 AT_ERROR("minlength should be >= 0");
28 }
29 if (self.dim() == 1 && self.numel() == 0) {
30 return at::zeros({minlength}, kLong);
31 }
32 if (self.dim() != 1 || *self.min().data_ptr<input_t>() < 0) {
33 AT_ERROR("bincount only supports 1-d non-negative integral inputs.");
34 }
35
36 bool has_weights = weights.defined();
37 if (has_weights && (weights.dim() != 1 || weights.size(0) != self.size(0))) {
38 AT_ERROR("weights should be 1-d and have the same length as input");
39 }
40
41 Tensor output;
42 int64_t self_size = self.size(0);
43 int64_t nbins = static_cast<int64_t>(*self.max().data_ptr<input_t>()) + 1L;
44 nbins = std::max(nbins, minlength); // at least minlength # of bins
45
46 const input_t* self_p = self.const_data_ptr<input_t>();
47 if (has_weights) {
48 output = at::zeros(
49 {nbins},
50 optTypeMetaToScalarType(weights.options().dtype_opt()),
51 weights.options().layout_opt(),
52 weights.options().device_opt(),
53 weights.options().pinned_memory_opt());
54 weights_t* output_p = output.data_ptr<weights_t>();
55 const weights_t* weights_p = weights.const_data_ptr<weights_t>();
56 for (const auto i : c10::irange(self_size)) {
57 output_p[self_p[i]] += weights_p[i];
58 }
59 } else {
60 output = at::zeros({nbins}, kLong);
61 int64_t* output_p = output.data_ptr<int64_t>();
62 for (const auto i : c10::irange(self_size)) {
63 output_p[self_p[i]] += 1L;
64 }
65 }
66 return output;
67 }
68 } // namespace
69
70 Tensor
_bincount_cpu(const Tensor & self,const std::optional<Tensor> & weights_opt,int64_t minlength)71 _bincount_cpu(const Tensor& self, const std::optional<Tensor>& weights_opt, int64_t minlength) {
72 // See [Note: hacky wrapper removal for optional tensor]
73 c10::MaybeOwned<Tensor> weights_maybe_owned = at::borrow_from_optional_tensor(weights_opt);
74 const Tensor& weights = *weights_maybe_owned;
75
76 return AT_DISPATCH_INTEGRAL_TYPES(self.scalar_type(), "bincount_cpu", [&] {
77 const auto scalar = weights.scalar_type();
78 if (scalar == ScalarType::Undefined || scalar == ScalarType::Float)
79 return _bincount_cpu_template<scalar_t, float>(self.contiguous(), weights.contiguous(), minlength);
80 return _bincount_cpu_template<scalar_t, double>(
81 self.contiguous(), weights.contiguous().to(kDouble), minlength);
82 });
83 }
84
85 } // namespace at::native
86