441 lines
19 KiB
C++
441 lines
19 KiB
C++
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
|
#include <ATen/core/Tensor.h>
|
|
#include <ATen/Dispatch.h>
|
|
|
|
#include <ATen/native/Histogram.h>
|
|
#include <ATen/native/Resize.h>
|
|
|
|
#ifndef AT_PER_OPERATOR_HEADERS
|
|
#include <ATen/Functions.h>
|
|
#include <ATen/NativeFunctions.h>
|
|
#else
|
|
#include <ATen/ops/_histogramdd_bin_edges.h>
|
|
#include <ATen/ops/_histogramdd_bin_edges_native.h>
|
|
#include <ATen/ops/_histogramdd_from_bin_cts.h>
|
|
#include <ATen/ops/_histogramdd_from_bin_cts_native.h>
|
|
#include <ATen/ops/_histogramdd_from_bin_tensors.h>
|
|
#include <ATen/ops/_histogramdd_from_bin_tensors_native.h>
|
|
#include <ATen/ops/aminmax.h>
|
|
#include <ATen/ops/empty.h>
|
|
#include <ATen/ops/histc_native.h>
|
|
#include <ATen/ops/histogram_native.h>
|
|
#include <ATen/ops/histogramdd_native.h>
|
|
#include <ATen/ops/linspace.h>
|
|
#endif
|
|
|
|
#include <numeric>
|
|
#include <tuple>
|
|
#include <vector>
|
|
#include <functional>
|
|
#include <c10/util/ArrayRef.h>
|
|
#include <c10/core/ScalarType.h>
|
|
#include <c10/core/DefaultDtype.h>
|
|
#include <c10/util/irange.h>
|
|
|
|
/* Implements a numpy-like histogramdd function running on cpu
|
|
* https://numpy.org/doc/stable/reference/generated/numpy.histogramdd.html
|
|
*
|
|
* See the docstr for torch.histogramdd in torch/functional.py for further explanation.
|
|
*
|
|
* - torch.histogramdd(input, bins, range=None, weight=None, density=False)
|
|
* input - tensor with shape (M, N). input is interpreted as M coordinates in N-dimensional space.
|
|
* If a tensor with more than 2 dimensions is passed, all but the last dimension will be flattened.
|
|
* bins - int[] of length N or tensor list of length N. If int[], defines the number of equal-width bins
|
|
* in each dimension. If tensor list, defines the sequences of bin edges, including rightmost edges,
|
|
* for each dimension.
|
|
* range - float[] of length 2 * N, optional. If specified, defines the leftmost and rightmost bin edges
|
|
* for each dimension.
|
|
* weight - tensor, optional. If provided, weight should have the same shape as input excluding its last dimension.
|
|
* Each N-dimensional value in input contributes its associated weight towards its bin's result.
|
|
* If weight is not specified, each value has weight 1 by default.
|
|
* density - bool, optional. If false (default), the result will contain the total count (weight) in each bin.
|
|
* If True, each count (weight) is divided by the total count (total weight), then divided by the
|
|
* volume of its associated bin.
|
|
*
|
|
* Returns:
|
|
* hist - N-dimensional tensor containing the values of the histogram.
|
|
* bin_edges - tensor list of length N containing the edges of the histogram bins in each dimension.
|
|
* Bins include their left edge and exclude their right edge, with the exception of the
|
|
* rightmost bin in each dimension which includes both of its edges.
|
|
*
|
|
* Restrictions are defined in histogram_check_inputs() and in select_outer_bin_edges().
|
|
*/
|
|
|
|
namespace at::native {
|
|
|
|
DEFINE_DISPATCH(histogramdd_stub);
|
|
DEFINE_DISPATCH(histogramdd_linear_stub);
|
|
DEFINE_DISPATCH(histogram_select_outer_bin_edges_stub);
|
|
|
|
namespace {
|
|
|
|
/* Checks properties of input tensors input, bins, and weight.
|
|
*/
|
|
void histogramdd_check_inputs(const Tensor& input, const TensorList& bins, const std::optional<Tensor>& weight) {
|
|
TORCH_CHECK(input.dim() >= 2, "torch.histogramdd: input tensor should have at least 2 dimensions, but got ",
|
|
input.dim());
|
|
|
|
const int64_t N = input.size(-1);
|
|
|
|
TORCH_CHECK(static_cast<int64_t>(bins.size()) == N, "torch.histogramdd: expected ", N, " sequences of bin edges for a ", N,
|
|
"-dimensional histogram but got ", bins.size());
|
|
|
|
auto input_dtype = input.dtype();
|
|
for (const auto dim : c10::irange(N)) {
|
|
const Tensor& dim_bins = bins[dim];
|
|
|
|
auto bins_dtype = dim_bins.dtype();
|
|
TORCH_CHECK(input_dtype == bins_dtype, "torch.histogramdd: input tensor and bins tensors should",
|
|
" have the same dtype, but got input with dtype ", input_dtype,
|
|
" and bins for dimension ", dim, " with dtype ", bins_dtype);
|
|
|
|
const int64_t dim_bins_dim = dim_bins.dim();
|
|
TORCH_CHECK(dim_bins_dim == 1, "torch.histogramdd: bins tensor should have one dimension,",
|
|
" but got ", dim_bins_dim, " dimensions in the bins tensor for dimension ", dim);
|
|
|
|
const int64_t numel = dim_bins.numel();
|
|
TORCH_CHECK(numel > 0, "torch.histogramdd: bins tensor should have at least 1 element,",
|
|
" but got ", numel, " elements in the bins tensor for dimension ", dim);
|
|
}
|
|
|
|
if (weight.has_value()) {
|
|
TORCH_CHECK(input.dtype() == weight.value().dtype(), "torch.histogramdd: if weight tensor is provided,"
|
|
" input tensor and weight tensor should have the same dtype, but got input(", input.dtype(), ")",
|
|
", and weight(", weight.value().dtype(), ")");
|
|
|
|
/* If a weight tensor is provided, we expect its shape to match that of
|
|
* the input tensor excluding its innermost dimension N.
|
|
*/
|
|
auto input_sizes = input.sizes().vec();
|
|
input_sizes.pop_back();
|
|
|
|
auto weight_sizes = weight.value().sizes().vec();
|
|
if (weight_sizes.empty()) {
|
|
// correctly handle scalars
|
|
weight_sizes = {1};
|
|
}
|
|
|
|
TORCH_CHECK(input_sizes == weight_sizes, "torch.histogramdd: if weight tensor is provided it should have"
|
|
" the same shape as the input tensor excluding its innermost dimension, but got input with shape ",
|
|
input.sizes(), " and weight with shape ", weight.value().sizes());
|
|
}
|
|
}
|
|
|
|
/* Checks properties of output tensors hist and bin_edges, then resizes them.
|
|
*/
|
|
void histogramdd_prepare_out(const Tensor& input, const std::vector<int64_t>& bin_ct,
|
|
const Tensor& hist, const TensorList& bin_edges) {
|
|
const int64_t N = input.size(-1);
|
|
|
|
TORCH_INTERNAL_ASSERT((int64_t)bin_ct.size() == N);
|
|
TORCH_INTERNAL_ASSERT((int64_t)bin_edges.size() == N);
|
|
|
|
TORCH_CHECK(input.dtype() == hist.dtype(), "torch.histogram: input tensor and hist tensor should",
|
|
" have the same dtype, but got input ", input.dtype(), " and hist ", hist.dtype());
|
|
|
|
for (const auto dim : c10::irange(N)) {
|
|
TORCH_CHECK(input.dtype() == bin_edges[dim].dtype(), "torch.histogram: input tensor and bin_edges tensor should",
|
|
" have the same dtype, but got input ", input.dtype(), " and bin_edges ", bin_edges[dim].dtype(),
|
|
" for dimension ", dim);
|
|
|
|
TORCH_CHECK(bin_ct[dim] > 0,
|
|
"torch.histogram(): bins must be > 0, but got ", bin_ct[dim], " for dimension ", dim);
|
|
|
|
at::native::resize_output(bin_edges[dim], bin_ct[dim] + 1);
|
|
}
|
|
|
|
at::native::resize_output(hist, bin_ct);
|
|
}
|
|
|
|
void histogramdd_prepare_out(const Tensor& input, TensorList bins,
|
|
const Tensor& hist, const TensorList& bin_edges) {
|
|
std::vector<int64_t> bin_ct(bins.size());
|
|
std::transform(bins.begin(), bins.end(), bin_ct.begin(), [](Tensor t) { return t.numel() - 1; });
|
|
histogramdd_prepare_out(input, bin_ct, hist, bin_edges);
|
|
}
|
|
|
|
/* Determines the outermost bin edges. For simplicity when calling into aminmax,
|
|
* assumes that input has already been reshaped to (M, N).
|
|
*/
|
|
std::pair<std::vector<double>, std::vector<double>>
|
|
select_outer_bin_edges(const Tensor& input, std::optional<c10::ArrayRef<double>> range) {
|
|
TORCH_INTERNAL_ASSERT(input.dim() == 2, "expected input to have shape (M, N)");
|
|
const int64_t N = input.size(-1);
|
|
|
|
// Default ranges for empty input matching numpy.histogram's default
|
|
std::vector<double> leftmost_edges(N, 0.);
|
|
std::vector<double> rightmost_edges(N, 1.);
|
|
|
|
if (range.has_value()) {
|
|
// range is specified
|
|
TORCH_CHECK((int64_t)range.value().size() == 2 * N, "torch.histogramdd: for a ", N, "-dimensional histogram",
|
|
" range should have ", 2 * N, " elements, but got ", range.value().size());
|
|
|
|
for (const auto dim : c10::irange(N)) {
|
|
leftmost_edges[dim] = range.value()[2 * dim];
|
|
rightmost_edges[dim] = range.value()[2 * dim + 1];
|
|
}
|
|
} else if (input.numel() > 0) {
|
|
// non-empty input
|
|
|
|
histogram_select_outer_bin_edges_stub(input.device().type(), input, N, leftmost_edges, rightmost_edges);
|
|
}
|
|
|
|
for (const auto dim : c10::irange(N)) {
|
|
double leftmost_edge = leftmost_edges[dim];
|
|
double rightmost_edge = rightmost_edges[dim];
|
|
|
|
TORCH_CHECK(std::isfinite(leftmost_edge) && std::isfinite(rightmost_edge),
|
|
"torch.histogramdd: dimension ", dim, "'s range [",
|
|
leftmost_edge, ", ", rightmost_edge, "] is not finite");
|
|
|
|
TORCH_CHECK(leftmost_edge <= rightmost_edge, "torch.histogramdd: min should not exceed max, but got",
|
|
" min ", leftmost_edge, " max ", rightmost_edge, " for dimension ", dim);
|
|
|
|
// Expand empty range to match numpy behavior and avoid division by 0 in normalization
|
|
if (leftmost_edge == rightmost_edge) {
|
|
leftmost_edges[dim] -= 0.5;
|
|
rightmost_edges[dim] += 0.5;
|
|
}
|
|
}
|
|
|
|
return std::make_pair(leftmost_edges, rightmost_edges);
|
|
}
|
|
|
|
/* histc's version of the logic for outermost bin edges.
|
|
*/
|
|
std::pair<double, double> histc_select_outer_bin_edges(const Tensor& input,
|
|
const Scalar& min, const Scalar& max) {
|
|
double leftmost_edge = min.to<double>();
|
|
double rightmost_edge = max.to<double>();
|
|
|
|
if (leftmost_edge == rightmost_edge && input.numel() > 0) {
|
|
auto extrema = aminmax(input);
|
|
leftmost_edge = std::get<0>(extrema).item<double>();
|
|
rightmost_edge = std::get<1>(extrema).item<double>();
|
|
}
|
|
|
|
if (leftmost_edge == rightmost_edge) {
|
|
leftmost_edge -= 1;
|
|
rightmost_edge += 1;
|
|
}
|
|
|
|
TORCH_CHECK(!(std::isinf(leftmost_edge) || std::isinf(rightmost_edge) ||
|
|
std::isnan(leftmost_edge) || std::isnan(rightmost_edge)),
|
|
"torch.histc: range of [", leftmost_edge, ", ", rightmost_edge, "] is not finite");
|
|
|
|
TORCH_CHECK(leftmost_edge < rightmost_edge, "torch.histc: max must be larger than min");
|
|
|
|
return std::make_pair(leftmost_edge, rightmost_edge);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
static std::vector<Tensor> allocate_bin_edges_tensors(const Tensor& self) {
|
|
TORCH_CHECK(self.dim() >= 2, "torch.histogramdd: input tensor should have at least 2 dimensions");
|
|
const int64_t N = self.size(-1);
|
|
std::vector<Tensor> bin_edges_out(N);
|
|
for (const auto dim : c10::irange(N)) {
|
|
bin_edges_out[dim] = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
}
|
|
return bin_edges_out;
|
|
}
|
|
|
|
/* Versions of histogramdd in which bins is a Tensor[] defining the sequences of bin edges.
|
|
*/
|
|
static Tensor& histogramdd_out(const Tensor& self, TensorList bins,
|
|
const std::optional<Tensor>& weight, bool density,
|
|
Tensor& hist, TensorList& bin_edges) {
|
|
histogramdd_check_inputs(self, bins, weight);
|
|
histogramdd_prepare_out(self, bins, hist, bin_edges);
|
|
|
|
for (const auto dim : c10::irange(bins.size())) {
|
|
bin_edges[dim].copy_(bins[dim]);
|
|
}
|
|
|
|
histogramdd_stub(self.device().type(), self, weight, density, hist, bin_edges);
|
|
return hist;
|
|
}
|
|
|
|
Tensor _histogramdd(const Tensor& self, TensorList bins,
|
|
const std::optional<Tensor>& weight, bool density) {
|
|
Tensor hist = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
std::vector<Tensor> bin_edges_out = allocate_bin_edges_tensors(self);
|
|
TensorList bin_edges_out_tl(bin_edges_out);
|
|
|
|
histogramdd_out(self, bins, weight, density, hist, bin_edges_out_tl);
|
|
return hist;
|
|
}
|
|
|
|
/* Versions of histogramdd in which bins is an int[]
|
|
* defining the number of bins in each dimension.
|
|
*/
|
|
static std::vector<Tensor>& histogramdd_bin_edges_out(const Tensor& self, IntArrayRef bin_ct,
|
|
std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density,
|
|
std::vector<Tensor>& bin_edges_out) {
|
|
TensorList bin_edges_out_tl(bin_edges_out);
|
|
|
|
const int64_t N = self.size(-1);
|
|
const int64_t M = std::accumulate(self.sizes().begin(), self.sizes().end() - 1,
|
|
(int64_t)1, std::multiplies<int64_t>());
|
|
Tensor reshaped_self = self.reshape({ M, N });
|
|
|
|
auto outer_bin_edges = select_outer_bin_edges(reshaped_self, range);
|
|
|
|
const int64_t bin_size = bin_ct.size();
|
|
TORCH_CHECK(
|
|
N == bin_size,
|
|
"histogramdd: The size of bins must be equal to the innermost dimension of the input.");
|
|
for (const auto dim : c10::irange(N)) {
|
|
at::linspace_out(bin_edges_out[dim], outer_bin_edges.first[dim], outer_bin_edges.second[dim],
|
|
bin_ct[dim] + 1);
|
|
}
|
|
|
|
return bin_edges_out;
|
|
}
|
|
|
|
std::vector<Tensor> histogramdd_bin_edges(const Tensor& self, IntArrayRef bin_ct,
|
|
std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density) {
|
|
std::vector<Tensor> bin_edges_out = allocate_bin_edges_tensors(self);
|
|
return histogramdd_bin_edges_out(self, bin_ct, range, weight, density, bin_edges_out);
|
|
}
|
|
|
|
static Tensor& histogramdd_out(const Tensor& self, IntArrayRef bin_ct,
|
|
std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density,
|
|
Tensor& hist, TensorList& bin_edges) {
|
|
std::vector<Tensor> bins = histogramdd_bin_edges(self, bin_ct, range, weight, density);
|
|
|
|
histogramdd_check_inputs(self, bins, weight);
|
|
histogramdd_prepare_out(self, bins, hist, bin_edges);
|
|
|
|
for (const auto dim : c10::irange(bins.size())) {
|
|
bin_edges[dim].copy_(bins[dim]);
|
|
}
|
|
|
|
histogramdd_linear_stub(self.device().type(), self, weight, density, hist, bin_edges, true);
|
|
return hist;
|
|
}
|
|
|
|
Tensor _histogramdd(const Tensor& self, IntArrayRef bin_ct,
|
|
std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density) {
|
|
Tensor hist = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
std::vector<Tensor> bin_edges_out = allocate_bin_edges_tensors(self);
|
|
TensorList bin_edges_out_tl(bin_edges_out);
|
|
|
|
histogramdd_out(self, bin_ct, range, weight, density, hist, bin_edges_out_tl);
|
|
return hist;
|
|
}
|
|
|
|
/* Versions of histogram in which bins is a Tensor defining the sequence of bin edges.
|
|
*/
|
|
std::tuple<Tensor&, Tensor&>
|
|
histogram_out(const Tensor& self, const Tensor& bins,
|
|
const std::optional<Tensor>& weight, bool density,
|
|
Tensor& hist, Tensor& bin_edges) {
|
|
Tensor reshaped_self = self.reshape({ self.numel(), 1 });
|
|
std::optional<Tensor> reshaped_weight = weight.has_value()
|
|
? weight.value().reshape({ weight.value().numel() }) : weight;
|
|
TensorList bins_in = bins;
|
|
TensorList bins_out = bin_edges;
|
|
|
|
histogramdd_out(reshaped_self, bins_in, reshaped_weight, density, hist, bins_out);
|
|
|
|
return std::forward_as_tuple(hist, bin_edges);
|
|
}
|
|
|
|
std::tuple<Tensor, Tensor>
|
|
histogram(const Tensor& self, const Tensor& bins,
|
|
const std::optional<Tensor>& weight, bool density) {
|
|
Tensor hist = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
Tensor bin_edges = at::empty({0}, bins.options(), MemoryFormat::Contiguous);
|
|
return histogram_out(self, bins, weight, density, hist, bin_edges);
|
|
}
|
|
|
|
/* Versions of histogram in which bins is an integer specifying the number of equal-width bins.
|
|
*/
|
|
std::tuple<Tensor&, Tensor&>
|
|
histogram_out(const Tensor& self, int64_t bin_ct, std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density,
|
|
Tensor& hist, Tensor& bin_edges) {
|
|
Tensor reshaped_self = self.reshape({ self.numel(), 1 });
|
|
std::optional<Tensor> reshaped_weight = weight.has_value()
|
|
? weight.value().reshape({ weight.value().numel() }) : weight;
|
|
TensorList bins_in = bin_edges;
|
|
TensorList bins_out = bin_edges;
|
|
|
|
histogramdd_prepare_out(reshaped_self, std::vector<int64_t>{bin_ct}, hist, bins_out);
|
|
auto outer_bin_edges = select_outer_bin_edges(reshaped_self, range);
|
|
at::linspace_out(bin_edges, outer_bin_edges.first[0], outer_bin_edges.second[0], bin_ct + 1);
|
|
|
|
histogramdd_check_inputs(reshaped_self, bins_in, reshaped_weight);
|
|
|
|
histogramdd_linear_stub(reshaped_self.device().type(), reshaped_self, reshaped_weight, density, hist, bin_edges, true);
|
|
return std::forward_as_tuple(hist, bin_edges);
|
|
}
|
|
|
|
std::tuple<Tensor, Tensor>
|
|
histogram(const Tensor& self, int64_t bin_ct, std::optional<c10::ArrayRef<double>> range,
|
|
const std::optional<Tensor>& weight, bool density) {
|
|
Tensor hist = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
Tensor bin_edges_out = at::empty({0}, self.options());
|
|
return histogram_out(self, bin_ct, range, weight, density, hist, bin_edges_out);
|
|
}
|
|
|
|
/* Narrowed interface for the legacy torch.histc function.
|
|
*/
|
|
Tensor& histogram_histc_out(const Tensor& self, int64_t bin_ct,
|
|
const Scalar& min, const Scalar& max, Tensor& hist) {
|
|
Tensor bin_edges = at::empty({0}, self.options());
|
|
|
|
Tensor reshaped = self.reshape({ self.numel(), 1 });
|
|
TensorList bins_in = bin_edges;
|
|
TensorList bins_out = bin_edges;
|
|
|
|
histogramdd_prepare_out(reshaped, std::vector<int64_t>{bin_ct}, hist, bins_out);
|
|
|
|
auto outer_bin_edges = histc_select_outer_bin_edges(self, min, max);
|
|
at::linspace_out(bin_edges, outer_bin_edges.first, outer_bin_edges.second, bin_ct + 1);
|
|
|
|
histogramdd_check_inputs(reshaped, bins_in, {});
|
|
|
|
histogramdd_linear_stub(reshaped.device().type(), reshaped,
|
|
std::optional<Tensor>(), false, hist, bin_edges, false);
|
|
return hist;
|
|
}
|
|
|
|
Tensor histogram_histc(const Tensor& self, int64_t bin_ct,
|
|
const Scalar& min, const Scalar& max) {
|
|
Tensor hist = at::empty({0}, self.options(), MemoryFormat::Contiguous);
|
|
return histogram_histc_out(self, bin_ct, min, max, hist);
|
|
}
|
|
|
|
std::tuple<Tensor, std::vector<Tensor>> histogramdd(
|
|
const Tensor &self, TensorList bins, std::optional<ArrayRef<double>> /*range*/,
|
|
const std::optional<Tensor> &weight, bool density) {
|
|
auto hist = at::_histogramdd_from_bin_tensors(self, bins, weight, density);
|
|
return std::tuple<Tensor, std::vector<Tensor>>{
|
|
std::move(hist), bins.vec()};
|
|
}
|
|
|
|
std::tuple<Tensor, std::vector<Tensor>> histogramdd(
|
|
const Tensor &self, IntArrayRef bins, std::optional<ArrayRef<double>> range,
|
|
const std::optional<Tensor> &weight, bool density) {
|
|
auto bin_edges = at::_histogramdd_bin_edges(self, bins, range, weight, density);
|
|
auto hist = at::_histogramdd_from_bin_cts(self, bins, range, weight, density);
|
|
return std::tuple<Tensor, std::vector<Tensor>>{
|
|
std::move(hist), std::move(bin_edges)};
|
|
}
|
|
|
|
std::tuple<Tensor, std::vector<Tensor>> histogramdd(
|
|
const Tensor &self, int64_t bins, std::optional<ArrayRef<double>> range,
|
|
const std::optional<Tensor> &weight, bool density) {
|
|
DimVector bins_v(self.size(-1), bins);
|
|
return at::native::histogramdd(self, bins_v, range, weight, density);
|
|
}
|
|
|
|
} // namespace at::native
|