Skip to content

[ET][Util] Add tensor utility functions #694

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[ET][Util] Add tensor utility functions #694

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions runtime/core/exec_aten/util/tensor_util.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -535,6 +535,18 @@ inline bool tensor_has_rank_greater_or_equal_to(
return true;
}

inline bool tensor_has_rank_smaller_or_equal_to(
exec_aten::Tensor t,
size_t rank) {
ET_LOG_MSG_AND_RETURN_IF_FALSE(
t.dim() <= rank,
"Expected tensor.dim() to be <= %zu, but got %zu",
static_cast<size_t>(rank),
static_cast<size_t>(t.dim()));

return true;
}

inline bool tensor_has_dim(exec_aten::Tensor t, int64_t d) {
if (t.dim() == 0) {
ET_LOG_MSG_AND_RETURN_IF_FALSE(
Expand All @@ -551,6 +563,25 @@ inline bool tensor_has_dim(exec_aten::Tensor t, int64_t d) {
return true;
}

inline bool tensor_dim_has_index(exec_aten::Tensor t, int64_t d, int64_t ix) {
// Indexing ops don't support zero-dim tensors
ET_CHECK(t.dim() != 0);
if (d < 0) {
d += t.dim();
}
// Dimension must have been already checked by tensor_has_dim
ET_CHECK(d >= 0 && d < t.dim());

ET_LOG_MSG_AND_RETURN_IF_FALSE(
ix >= -t.size(d) && ix < t.size(d),
"index %" PRId64 " out of range [-%zu,%zu) at dimension %" PRId64 ")",
ix,
static_cast<size_t>(t.size(d)),
static_cast<size_t>(t.size(d)),
d);
return true;
}

inline bool tensors_have_same_size_at_dims(
exec_aten::Tensor a,
size_t dim_a,
Expand Down