Spaces:
Running
Running
namespace at::native { | |
// original values given by raw_*. If an original value is not contiguous, will make a contiguous copy to | |
// the corresponding trimmed_* value. Additionally, if the dtypes of the boundary and input tensor do not | |
// match, will change them to be a common super type so comparisons are done between the same types. | |
// For any trimmed_* tensor, if its outgoing value matches what it was incoming (typically null), then the | |
// corresponding raw_* version should be used since it was already contiguous of the right type. | |
inline void searchsorted_maybe_trim_input_tensors( | |
Tensor& trimmed_input, | |
Tensor& trimmed_boundaries, | |
Tensor& trimmed_sorter, | |
const Tensor& raw_input, | |
const Tensor& raw_boundaries, | |
const Tensor& raw_sorter) { | |
bool in_is_contiguous = raw_input.is_contiguous(); | |
bool bd_is_contiguous = raw_boundaries.is_contiguous(); | |
bool sort_is_contiguous = raw_sorter.is_contiguous(); | |
if (!in_is_contiguous) { | |
TORCH_WARN_ONCE("torch.searchsorted(): input value tensor is non-contiguous, this will lower the performance due " | |
"to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous input value " | |
"tensor if possible. This message will only appear once per program."); | |
trimmed_input = raw_input.contiguous(); | |
} | |
if (!bd_is_contiguous) { | |
TORCH_WARN_ONCE("torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due " | |
"to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary " | |
"tensor if possible. This message will only appear once per program."); | |
trimmed_boundaries = raw_boundaries.contiguous(); | |
} | |
if (!sort_is_contiguous) { | |
TORCH_WARN_ONCE("torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due " | |
"to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter " | |
"tensor if possible. This message will only appear once per program."); | |
trimmed_sorter = raw_sorter.contiguous(); | |
} | |
if (raw_input.dtype() != raw_boundaries.dtype()) { | |
at::native::ResultTypeState state = {}; | |
state = at::native::update_result_type_state(raw_boundaries, state); | |
state = at::native::update_result_type_state(raw_input, state); | |
ScalarType common_stype = at::native::result_type(state); | |
TORCH_INTERNAL_ASSERT(common_stype != ScalarType::Undefined); | |
if (common_stype != raw_input.scalar_type()) { | |
trimmed_input = in_is_contiguous ? raw_input.to(common_stype) : trimmed_input.to(common_stype); | |
} | |
if (common_stype != raw_boundaries.scalar_type()) { | |
trimmed_boundaries = bd_is_contiguous ? raw_boundaries.to(common_stype) : trimmed_boundaries.to(common_stype); | |
} | |
} | |
} | |
/* unused but needed for internal jagged tensor class */ | |
inline void searchsorted_maybe_trim_input_tensors( | |
Tensor& trimmed_input, | |
Tensor& trimmed_boundaries, | |
const Tensor& raw_input, | |
const Tensor& raw_boundaries) { | |
Tensor trimmed_sorter; | |
Tensor raw_sorter; | |
return searchsorted_maybe_trim_input_tensors( | |
trimmed_input, | |
trimmed_boundaries, | |
trimmed_sorter, | |
raw_input, | |
raw_boundaries, | |
raw_sorter); | |
} | |
inline bool searchsorted_dims_matched_before_last_dim(const Tensor& boundaries, const Tensor& input) { | |
if (boundaries.dim() != input.dim()) { | |
return false; | |
} | |
const auto& dims_bd = boundaries.sizes(); | |
const auto& dims_in = input.sizes(); | |
for (int64_t dim = 0; dim + 1 < boundaries.dim(); ++dim) { | |
if (dims_bd[dim] != dims_in[dim]) { | |
return false; | |
} | |
} | |
return true; | |
} | |
inline Tensor searchsorted_scalar_tensor(const Scalar& scalar, const c10::Device& device) { | |
auto tensor = c10::scalar_to_tensor(scalar, device); | |
// This is to adopt the scalar promotion rules defined in native/TypeProperties.h | |
// So we have the same type promotion rules as binary operations. | |
tensor.unsafeGetTensorImpl()->set_wrapped_number(true); | |
return tensor; | |
} | |
inline void searchsorted_pre_check( | |
const Tensor& boundaries, | |
const Tensor& input, | |
const Tensor& output, | |
const bool out_int32, | |
const bool right, | |
const c10::optional<c10::string_view> side_opt, | |
const Tensor& sorter) { | |
if (side_opt) { | |
const c10::string_view side = *side_opt; | |
TORCH_CHECK(side == "left" || side == "right", "torch.searchsorted(): side can only be 'left' or 'right' but ", | |
"got ", side); | |
// assume the user has not explicitly set (right=False, side="right") | |
TORCH_CHECK(!right || side == "right", "torch.searchsorted(): side and right can't be set to opposites, got side " | |
"of ", side, " while right was True"); | |
} | |
TORCH_CHECK(boundaries.device() == input.device(), "torch.searchsorted(): boundaries and input value tensors ", | |
"should have same device type, but got boundaries tensor device type ", boundaries.device(), " and input value ", | |
"tensor device type ", input.device()); | |
if (sorter.defined()) { | |
TORCH_CHECK(sorter.device() == boundaries.device(), "torch.searchsorted(): sorter and boundary tensors should ", | |
"have same device type, but got sorter tensor device type ", sorter.device(), " and input value tensor ", | |
"device type ", boundaries.device()); | |
TORCH_CHECK(sorter.sizes() == boundaries.sizes(), "torch.searchsorted(): boundary and sorter must have the same " | |
"size, but got boundary tensor ", boundaries.sizes(), "and got sorter tensor ", sorter.sizes()); | |
TORCH_CHECK(sorter.scalar_type() == ScalarType::Long, "torch.searchsorted(): sorter must be a tensor of long ", | |
"dtype but got dtype ", sorter.scalar_type()); | |
if (sorter.numel() > 0) { | |
auto minmax = sorter.aminmax(); | |
int64_t vmin = std::get<0>(minmax).item().toLong(); | |
int64_t vmax = std::get<1>(minmax).item().toLong(); | |
TORCH_CHECK(vmin >= 0 && vmax < sorter.sizes().back(), "torch.searchsorted(): sorter index out of range"); | |
} | |
} | |
TORCH_CHECK(input.dim() > 0 || (input.dim() == 0 && input.numel() == 1 && boundaries.dim() == 1), | |
"torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we got ", | |
"boundaries tensor dim(", boundaries.dim(), ") and input value's dim(", input.dim(), ") numel(", | |
input.numel(), ")"); | |
TORCH_CHECK(boundaries.dim() != 0, "torch.searchsorted(): boundaries tensor should have positive dimension, but ", | |
"got 0 dimension"); | |
TORCH_CHECK(boundaries.dim() == 1 || searchsorted_dims_matched_before_last_dim(boundaries, input), | |
"torch.searchsorted(): boundaries tensor should be 1 dimension or the first N-1 dimensions of boundaries tensor ", | |
"and input value tensor must match, but we got boundaries tensor ", boundaries.sizes(), " and input value tensor ", | |
input.sizes()); | |
ScalarType output_dtype = output.scalar_type(); | |
TORCH_CHECK( | |
(output_dtype == ScalarType::Long && !out_int32) || | |
(output_dtype == ScalarType::Int && out_int32), | |
"torch.searchsorted(): output tensor's dtype is wrong, it can only be Int(int32) or Long(int64) depending on ", | |
"whether out_int32 flag is True, but we got output tensor's dtype ", output_dtype, | |
" and out_int32 flag is ", (out_int32 ? "True" : "False")); | |
if (out_int32) { | |
TORCH_CHECK(boundaries.sizes().back() < INT_MAX, | |
"torch.searchsorted(): the size of boundaries' last dimension should be less than ", INT_MAX, ", but we got ", | |
boundaries.sizes().back()); | |
} | |
} | |
} // namespace at::native | |