Skip to content
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

Jump to bottom

scalar support to more binary ops #2093

Merged
merged 2 commits into from
Jul 21, 2020
Merged

scalar support to more binary ops #2093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1e9273f
scalar support for more binary ops
yxlao Jul 21, 2020
520ca90
fix boolean advanced indexing
yxlao Jul 21, 2020
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
72 changes: 72 additions & 0 deletions cpp/open3d/core/Tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -594,6 +594,10 @@ class Tensor {
/// non-zero values will be treated as True.
Tensor LogicalAnd(const Tensor& value) const;
Tensor operator&&(const Tensor& value) const { return LogicalAnd(value); }
template <typename T>
Tensor LogicalAnd(T scalar_value) const {
return LogicalAnd(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise logical and of tensors, in-place. This operation won't
/// change the tensor's dtype.
Expand All @@ -602,13 +606,21 @@ class Tensor {
/// will be treated as True. The tensor will be filled with 0 or 1 casted to
/// the tensor's dtype.
Tensor LogicalAnd_(const Tensor& value);
template <typename T>
Tensor LogicalAnd_(T scalar_value) {
return LogicalAnd_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise logical or of tensors, returning a new boolean tensor.
///
/// If the tensor is not boolean, zero will be treated as False, while
/// non-zero values will be treated as True.
Tensor LogicalOr(const Tensor& value) const;
Tensor operator||(const Tensor& value) const { return LogicalOr(value); }
template <typename T>
Tensor LogicalOr(T scalar_value) const {
return LogicalOr(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise logical or of tensors, in-place. This operation won't
/// change the tensor's dtype.
Expand All @@ -617,13 +629,21 @@ class Tensor {
/// will be treated as True. The tensor will be filled with 0 or 1 casted to
/// the tensor's dtype.
Tensor LogicalOr_(const Tensor& value);
template <typename T>
Tensor LogicalOr_(T scalar_value) {
return LogicalOr_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise logical exclusive-or of tensors, returning a new boolean
/// tensor.
///
/// If the tensor is not boolean, zero will be treated as False, while
/// non-zero values will be treated as True.
Tensor LogicalXor(const Tensor& value) const;
template <typename T>
Tensor LogicalXor(T scalar_value) const {
return LogicalXor(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise logical exclusive-or of tensors, in-place. This operation
/// won't change the tensor's dtype.
Expand All @@ -632,56 +652,108 @@ class Tensor {
/// non-zero values will be treated as True. The tensor will be filled with
/// 0 or 1 casted to the tensor's dtype.
Tensor LogicalXor_(const Tensor& value);
template <typename T>
Tensor LogicalXor_(T scalar_value) {
return LogicalXor_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise greater-than of tensors, returning a new boolean tensor.
Tensor Gt(const Tensor& value) const;
Tensor operator>(const Tensor& value) const { return Gt(value); }
template <typename T>
Tensor Gt(T scalar_value) const {
return Gt(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise greater-than of tensors, in-place. This operation
/// won't change the tensor's dtype.
Tensor Gt_(const Tensor& value);
template <typename T>
Tensor Gt_(T scalar_value) {
return Gt_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise less-than of tensors, returning a new boolean tensor.
Tensor Lt(const Tensor& value) const;
Tensor operator<(const Tensor& value) const { return Lt(value); }
template <typename T>
Tensor Lt(T scalar_value) const {
return Lt(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise less-than of tensors, in-place. This operation won't change
/// the tensor's dtype.
Tensor Lt_(const Tensor& value);
template <typename T>
Tensor Lt_(T scalar_value) {
return Lt_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise greater-than-or-equals-to of tensors, returning a new
/// boolean tensor.
Tensor Ge(const Tensor& value) const;
Tensor operator>=(const Tensor& value) const { return Ge(value); }
template <typename T>
Tensor Ge(T scalar_value) const {
return Ge(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise greater-than-or-equals-to of tensors, in-place. This
/// operation won't change the tensor's dtype.
Tensor Ge_(const Tensor& value);
template <typename T>
Tensor Ge_(T scalar_value) {
return Ge_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise less-than-or-equals-to of tensors, returning a new boolean
/// tensor.
Tensor Le(const Tensor& value) const;
Tensor operator<=(const Tensor& value) const { return Le(value); }
template <typename T>
Tensor Le(T scalar_value) const {
return Le(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise less-than-or-equals-to of tensors, in-place. This operation
/// won't change the tensor's dtype.
Tensor Le_(const Tensor& value);
template <typename T>
Tensor Le_(T scalar_value) {
return Le_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise equals-to of tensors, returning a new boolean tensor.
Tensor Eq(const Tensor& value) const;
Tensor operator==(const Tensor& value) const { return Eq(value); }
template <typename T>
Tensor Eq(T scalar_value) const {
return Eq(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise equals-to of tensors, in-place. This
/// operation won't change the tensor's dtype.
Tensor Eq_(const Tensor& value);
template <typename T>
Tensor Eq_(T scalar_value) {
return Eq_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise not-equals-to of tensors, returning a new boolean tensor.
Tensor Ne(const Tensor& value) const;
Tensor operator!=(const Tensor& value) const { return Ne(value); }
template <typename T>
Tensor Ne(T scalar_value) const {
return Ne(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Element-wise equals-to of tensors, in-place. This
/// operation won't change the tensor's dtype.
Tensor Ne_(const Tensor& value);
template <typename T>
Tensor Ne_(T scalar_value) {
return Ne_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
}

/// Find the indices of the elements that are non-zero. Returns a vector of
/// int64 Tensors, each containing the indices of the non-zero elements in
Expand Down
2 changes: 1 addition & 1 deletion cpp/open3d/core/TensorKey.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ class TensorKey {
static TensorKey Slice(NoneType start, NoneType stop, int64_t step);
static TensorKey Slice(NoneType start, NoneType stop, NoneType step);

/// Construct an TensorKeyMode::IndexTensor type TensorKey (advnced
/// Construct an TensorKeyMode::IndexTensor type TensorKey (advanced
/// indexing).
static TensorKey IndexTensor(const Tensor& index_tensor);

Expand Down
133 changes: 40 additions & 93 deletions cpp/pybind/core/tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,20 @@
#include "open3d/core/Tensor.h"
#include "open3d/core/TensorKey.h"

#define CONST_ARG const
#define NON_CONST_ARG
#define BIND_BINARY_OP_ALL_DTYPES(py_name, cpp_name, self_const) \
tensor.def(#py_name, \
[](self_const core::Tensor& self, const core::Tensor& other) { \
return self.cpp_name(other); \
}); \
tensor.def(#py_name, &core::Tensor::cpp_name<float>); \
tensor.def(#py_name, &core::Tensor::cpp_name<double>); \
tensor.def(#py_name, &core::Tensor::cpp_name<int32_t>); \
tensor.def(#py_name, &core::Tensor::cpp_name<int64_t>); \
tensor.def(#py_name, &core::Tensor::cpp_name<uint8_t>); \
tensor.def(#py_name, &core::Tensor::cpp_name<bool>);

namespace open3d {

template <typename T>
Expand Down Expand Up @@ -262,101 +276,34 @@ void pybind_core_tensor(py::module& m) {
tensor.def("to", &core::Tensor::To);

// Binary element-wise ops
tensor.def("add", [](const core::Tensor& self, const core::Tensor& other) {
return self.Add(other);
});
tensor.def("add", &core::Tensor::Add<float>);
tensor.def("add", &core::Tensor::Add<double>);
tensor.def("add", &core::Tensor::Add<int32_t>);
tensor.def("add", &core::Tensor::Add<int64_t>);
tensor.def("add", &core::Tensor::Add<uint8_t>);
tensor.def("add", &core::Tensor::Add<bool>);
tensor.def("add_", [](core::Tensor& self, const core::Tensor& other) {
return self.Add_(other);
});
tensor.def("add_", &core::Tensor::Add_<float>);
tensor.def("add_", &core::Tensor::Add_<double>);
tensor.def("add_", &core::Tensor::Add_<int32_t>);
tensor.def("add_", &core::Tensor::Add_<int64_t>);
tensor.def("add_", &core::Tensor::Add_<uint8_t>);
tensor.def("add_", &core::Tensor::Add_<bool>);

tensor.def("sub", [](const core::Tensor& self, const core::Tensor& other) {
return self.Sub(other);
});
tensor.def("sub", &core::Tensor::Sub<float>);
tensor.def("sub", &core::Tensor::Sub<double>);
tensor.def("sub", &core::Tensor::Sub<int32_t>);
tensor.def("sub", &core::Tensor::Sub<int64_t>);
tensor.def("sub", &core::Tensor::Sub<uint8_t>);
tensor.def("sub", &core::Tensor::Sub<bool>);
tensor.def("sub_", [](core::Tensor& self, const core::Tensor& other) {
return self.Sub_(other);
});
tensor.def("sub_", &core::Tensor::Sub_<float>);
tensor.def("sub_", &core::Tensor::Sub_<double>);
tensor.def("sub_", &core::Tensor::Sub_<int32_t>);
tensor.def("sub_", &core::Tensor::Sub_<int64_t>);
tensor.def("sub_", &core::Tensor::Sub_<uint8_t>);
tensor.def("sub_", &core::Tensor::Sub_<bool>);

tensor.def("mul", [](const core::Tensor& self, const core::Tensor& other) {
return self.Mul(other);
});
tensor.def("mul", &core::Tensor::Mul<float>);
tensor.def("mul", &core::Tensor::Mul<double>);
tensor.def("mul", &core::Tensor::Mul<int32_t>);
tensor.def("mul", &core::Tensor::Mul<int64_t>);
tensor.def("mul", &core::Tensor::Mul<uint8_t>);
tensor.def("mul", &core::Tensor::Mul<bool>);
tensor.def("mul_", [](core::Tensor& self, const core::Tensor& other) {
return self.Mul_(other);
});
tensor.def("mul_", &core::Tensor::Mul_<float>);
tensor.def("mul_", &core::Tensor::Mul_<double>);
tensor.def("mul_", &core::Tensor::Mul_<int32_t>);
tensor.def("mul_", &core::Tensor::Mul_<int64_t>);
tensor.def("mul_", &core::Tensor::Mul_<uint8_t>);
tensor.def("mul_", &core::Tensor::Mul_<bool>);

tensor.def("div", [](const core::Tensor& self, const core::Tensor& other) {
return self.Div(other);
});
tensor.def("div", &core::Tensor::Div<float>);
tensor.def("div", &core::Tensor::Div<double>);
tensor.def("div", &core::Tensor::Div<int32_t>);
tensor.def("div", &core::Tensor::Div<int64_t>);
tensor.def("div", &core::Tensor::Div<uint8_t>);
tensor.def("div", &core::Tensor::Div<bool>);
tensor.def("div_", [](core::Tensor& self, const core::Tensor& other) {
return self.Div_(other);
});
tensor.def("div_", &core::Tensor::Div_<float>);
tensor.def("div_", &core::Tensor::Div_<double>);
tensor.def("div_", &core::Tensor::Div_<int32_t>);
tensor.def("div_", &core::Tensor::Div_<int64_t>);
tensor.def("div_", &core::Tensor::Div_<uint8_t>);
tensor.def("div_", &core::Tensor::Div_<bool>);
BIND_BINARY_OP_ALL_DTYPES(add, Add, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(add_, Add_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(sub, Sub, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(sub_, Sub_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(mul, Mul, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(mul_, Mul_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(div, Div, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(div_, Div_, NON_CONST_ARG);

// Binary boolean element-wise ops
tensor.def("logical_and", &core::Tensor::LogicalAnd);
tensor.def("logical_and_", &core::Tensor::LogicalAnd_);
tensor.def("logical_or", &core::Tensor::LogicalOr);
tensor.def("logical_or_", &core::Tensor::LogicalOr_);
tensor.def("logical_xor", &core::Tensor::LogicalXor);
tensor.def("logical_xor_", &core::Tensor::LogicalXor_);
tensor.def("gt", &core::Tensor::Gt);
tensor.def("gt_", &core::Tensor::Gt_);
tensor.def("lt", &core::Tensor::Lt);
tensor.def("lt_", &core::Tensor::Lt_);
tensor.def("ge", &core::Tensor::Ge);
tensor.def("ge_", &core::Tensor::Ge_);
tensor.def("le", &core::Tensor::Le);
tensor.def("le_", &core::Tensor::Le_);
tensor.def("eq", &core::Tensor::Eq);
tensor.def("eq_", &core::Tensor::Eq_);
tensor.def("ne", &core::Tensor::Ne);
tensor.def("ne_", &core::Tensor::Ne_);
BIND_BINARY_OP_ALL_DTYPES(logical_and, LogicalAnd, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(logical_and_, LogicalAnd_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(logical_or, LogicalOr, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(logical_or_, LogicalOr_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(logical_xor, LogicalXor, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(logical_xor_, LogicalXor_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(gt, Gt, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(gt_, Gt_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(lt, Lt, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(lt_, Lt_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(ge, Ge, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(ge_, Ge_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(le, Le, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(le_, Le_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(eq, Eq, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(eq_, Eq_, NON_CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(ne, Ne, CONST_ARG);
BIND_BINARY_OP_ALL_DTYPES(ne_, Ne_, NON_CONST_ARG);

// Getters and setters as peoperty
tensor.def_property_readonly("shape", [](const core::Tensor& tensor) {
Expand Down
4 changes: 3 additions & 1 deletion python/open3d/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ class Tensor(o3d.pybind.core.Tensor):
"""

def __init__(self, data, dtype=None, device=None):
if isinstance(data, tuple) or isinstance(data, list):
if isinstance(data, (tuple, list, int, float)):
data = np.array(data)
if not isinstance(data, np.ndarray):
raise ValueError("data must be a list, tuple, or Numpy array.")
Expand All @@ -126,6 +126,8 @@ def __getitem__(self, key):

@cast_to_py_tensor
def __setitem__(self, key, value):
if not isinstance(value, Tensor):
value = Tensor(value, self.dtype, self.device)
if isinstance(key, tuple):
o3d_tensor_keys = [_to_o3d_tensor_key(k) for k in key]
super(Tensor, self)._setitem_vector(o3d_tensor_keys, value)
Expand Down
Loading