Tensor.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#pragma once

#include <cassert>
#include <cstddef>
#include <memory>
#include <string>
#include <type_traits>

#include "open3d/core/Blob.h"
#include "open3d/core/DLPack.h"
#include "open3d/core/Device.h"
#include "open3d/core/Dtype.h"
#include "open3d/core/ShapeUtil.h"
#include "open3d/core/SizeVector.h"
#include "open3d/core/TensorKey.h"

namespace open3d {
namespace core {

class Tensor {
public:
    Tensor(){};

    Tensor(const SizeVector& shape,
           Dtype dtype,
           const Device& device = Device("CPU:0"))
        : shape_(shape),
          strides_(DefaultStrides(shape)),
          dtype_(dtype),
          blob_(std::make_shared<Blob>(shape.NumElements() * dtype.ByteSize(),
                                       device)) {
        data_ptr_ = blob_->GetDataPtr();
    }

    template <typename T>
    Tensor(const std::vector<T>& init_vals,
           const SizeVector& shape,
           Dtype dtype,
           const Device& device = Device("CPU:0"))
        : Tensor(shape, dtype, device) {
        // Check number of elements
        if (static_cast<int64_t>(init_vals.size()) != shape_.NumElements()) {
            utility::LogError(
                    "Tensor initialization values' size {} does not match the "
                    "shape {}",
                    init_vals.size(), shape_.NumElements());
        }

        // Check data types
        AssertTemplateDtype<T>();
        if (!std::is_pod<T>()) {
            utility::LogError("Object must be a POD.");
        }

        // Copy data to blob
        MemoryManager::MemcpyFromHost(blob_->GetDataPtr(), GetDevice(),
                                      init_vals.data(),
                                      init_vals.size() * dtype.ByteSize());
    }

    template <typename T>
    Tensor(const T* init_vals,
           const SizeVector& shape,
           Dtype dtype,
           const Device& device = Device("CPU:0"))
        : Tensor(shape, dtype, device) {
        // Check data types
        AssertTemplateDtype<T>();

        // Copy data to blob
        MemoryManager::MemcpyFromHost(blob_->GetDataPtr(), GetDevice(),
                                      init_vals,
                                      shape_.NumElements() * dtype.ByteSize());
    }

    Tensor(const SizeVector& shape,
           const SizeVector& strides,
           void* data_ptr,
           Dtype dtype,
           const std::shared_ptr<Blob>& blob)
        : shape_(shape),
          strides_(strides),
          data_ptr_(data_ptr),
          dtype_(dtype),
          blob_(blob) {}

    Tensor(const Tensor& other) = default;

    Tensor(Tensor&& other) = default;

    Tensor& operator=(const Tensor& other) &;

    Tensor& operator=(Tensor&& other) &;

    Tensor& operator=(const Tensor& other) &&;

    Tensor& operator=(Tensor&& other) &&;

    template <typename T>
    Tensor& operator=(const T& v) && {
        if (shape_.size() != 0) {
            utility::LogError(
                    "Assignment with scalar only works for scalar Tensor of "
                    "shape ()");
        }
        DISPATCH_DTYPE_TO_TEMPLATE_WITH_BOOL(GetDtype(), [&]() {
            scalar_t casted_v = static_cast<scalar_t>(v);
            MemoryManager::MemcpyFromHost(GetDataPtr(), GetDevice(), &casted_v,
                                          sizeof(scalar_t));
        });
        return *this;
    }

    template <typename Object>
    Tensor& AssignObject(const Object& v) && {
        if (shape_.size() != 0) {
            utility::LogError(
                    "Assignment with scalar only works for scalar Tensor of "
                    "shape ()");
        }
        AssertTemplateDtype<Object>();
        MemoryManager::MemcpyFromHost(GetDataPtr(), GetDevice(), &v,
                                      sizeof(Object));
        return *this;
    }

    template <typename Scalar>
    void Fill(Scalar v);

    template <typename Object>
    void FillObject(const Object& v);

    static Tensor Empty(const SizeVector& shape,
                        Dtype dtype,
                        const Device& device = Device("CPU:0"));

    static Tensor EmptyLike(const Tensor& other) {
        return Tensor::Empty(other.shape_, other.dtype_, other.GetDevice());
    }

    template <typename T>
    static Tensor Full(const SizeVector& shape,
                       T fill_value,
                       Dtype dtype,
                       const Device& device = Device("CPU:0")) {
        Tensor t = Empty(shape, dtype, device);
        t.Fill(fill_value);
        return t;
    }

    static Tensor Zeros(const SizeVector& shape,
                        Dtype dtype,
                        const Device& device = Device("CPU:0"));

    static Tensor Ones(const SizeVector& shape,
                       Dtype dtype,
                       const Device& device = Device("CPU:0"));

    static Tensor Eye(int64_t n, Dtype dtype, const Device& device);

    static Tensor Diag(const Tensor& input);

    Tensor GetItem(const TensorKey& tk) const;

    Tensor GetItem(const std::vector<TensorKey>& tks) const;

    Tensor SetItem(const Tensor& value);

    Tensor SetItem(const TensorKey& tk, const Tensor& value);

    Tensor SetItem(const std::vector<TensorKey>& tks, const Tensor& value);

    void Assign(const Tensor& other);

    Tensor Broadcast(const SizeVector& dst_shape) const;

    Tensor Expand(const SizeVector& dst_shape) const;

    Tensor Reshape(const SizeVector& dst_shape) const;

    Tensor View(const SizeVector& dst_shape) const;

    Tensor Copy(const Device& device) const;

    Tensor Copy() const { return Copy(GetDevice()); };

    void CopyFrom(const Tensor& other);

    void ShallowCopyFrom(const Tensor& other);

    Tensor To(Dtype dtype, bool copy = false) const;

    std::string ToString(bool with_suffix = true,
                         const std::string& indent = "") const;

    Tensor operator[](int64_t i) const;

    Tensor IndexExtract(int64_t dim, int64_t idx) const;

    Tensor Slice(int64_t dim,
                 int64_t start,
                 int64_t stop,
                 int64_t step = 1) const;

    Tensor AsRvalue() const { return *this; }

    Tensor IndexGet(const std::vector<Tensor>& index_tensors) const;

    void IndexSet(const std::vector<Tensor>& index_tensors,
                  const Tensor& src_tensor);

    Tensor Permute(const SizeVector& dims) const;

    Tensor AsStrided(const SizeVector& new_shape,
                     const SizeVector& new_strides) const;

    Tensor Transpose(int64_t dim0, int64_t dim1) const;

    Tensor T() const;

    template <typename T>
    T Item() const {
        if (shape_.size() != 0) {
            utility::LogError("Item only works for scalar Tensor of shape ()");
        }
        AssertTemplateDtype<T>();
        T value;
        MemoryManager::MemcpyToHost(&value, data_ptr_, GetDevice(), sizeof(T));
        return value;
    }

    Tensor Add(const Tensor& value) const;
    template <typename T>
    Tensor Add(T scalar_value) const {
        return Add(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator+(const Tensor& value) const { return Add(value); }
    template <typename T>
    Tensor operator+(T scalar_value) const {
        return Add(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Add_(const Tensor& value);
    template <typename T>
    Tensor Add_(T scalar_value) {
        return Add_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator+=(const Tensor& value) { return Add_(value); }
    template <typename T>
    Tensor operator+=(T scalar_value) {
        return Add_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Sub(const Tensor& value) const;
    template <typename T>
    Tensor Sub(T scalar_value) const {
        return Sub(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator-(const Tensor& value) const { return Sub(value); }
    template <typename T>
    Tensor operator-(T scalar_value) const {
        return Sub(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Sub_(const Tensor& value);
    template <typename T>
    Tensor Sub_(T scalar_value) {
        return Sub_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator-=(const Tensor& value) { return Sub_(value); }
    template <typename T>
    Tensor operator-=(T scalar_value) {
        return Sub_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Mul(const Tensor& value) const;
    template <typename T>
    Tensor Mul(T scalar_value) const {
        return Mul(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator*(const Tensor& value) const { return Mul(value); }
    template <typename T>
    Tensor operator*(T scalar_value) const {
        return Mul(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Mul_(const Tensor& value);
    template <typename T>
    Tensor Mul_(T scalar_value) {
        return Mul_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator*=(const Tensor& value) { return Mul_(value); }
    template <typename T>
    Tensor operator*=(T scalar_value) {
        return Mul_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Div(const Tensor& value) const;
    template <typename T>
    Tensor Div(T scalar_value) const {
        return Div(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator/(const Tensor& value) const { return Div(value); }
    template <typename T>
    Tensor operator/(T scalar_value) const {
        return Div(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Div_(const Tensor& value);
    template <typename T>
    Tensor Div_(T scalar_value) {
        return Div_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }
    Tensor operator/=(const Tensor& value) { return Div_(value); }
    template <typename T>
    Tensor operator/=(T scalar_value) {
        return Div_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor Sum(const SizeVector& dims, bool keepdim = false) const;

    Tensor Mean(const SizeVector& dims, bool keepdim = false) const;

    Tensor Prod(const SizeVector& dims, bool keepdim = false) const;

    Tensor Min(const SizeVector& dims, bool keepdim = false) const;

    Tensor Max(const SizeVector& dims, bool keepdim = false) const;

    Tensor ArgMin(const SizeVector& dims) const;

    Tensor ArgMax(const SizeVector& dims) const;

    Tensor Sqrt() const;

    Tensor Sqrt_();

    Tensor Sin() const;

    Tensor Sin_();

    Tensor Cos() const;

    Tensor Cos_();

    Tensor Neg() const;

    Tensor Neg_();

    Tensor Exp() const;

    Tensor Exp_();

    Tensor Abs() const;

    Tensor Abs_();
    Tensor LogicalNot() const;

    Tensor LogicalNot_();

    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()));
    }

    Tensor LogicalAnd_(const Tensor& value);
    template <typename T>
    Tensor LogicalAnd_(T scalar_value) {
        return LogicalAnd_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor LogicalOr_(const Tensor& value);
    template <typename T>
    Tensor LogicalOr_(T scalar_value) {
        return LogicalOr_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor LogicalXor(const Tensor& value) const;
    template <typename T>
    Tensor LogicalXor(T scalar_value) const {
        return LogicalXor(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    Tensor LogicalXor_(const Tensor& value);
    template <typename T>
    Tensor LogicalXor_(T scalar_value) {
        return LogicalXor_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Gt_(const Tensor& value);
    template <typename T>
    Tensor Gt_(T scalar_value) {
        return Gt_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Lt_(const Tensor& value);
    template <typename T>
    Tensor Lt_(T scalar_value) {
        return Lt_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Ge_(const Tensor& value);
    template <typename T>
    Tensor Ge_(T scalar_value) {
        return Ge_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Le_(const Tensor& value);
    template <typename T>
    Tensor Le_(T scalar_value) {
        return Le_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Eq_(const Tensor& value);
    template <typename T>
    Tensor Eq_(T scalar_value) {
        return Eq_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    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()));
    }

    Tensor Ne_(const Tensor& value);
    template <typename T>
    Tensor Ne_(T scalar_value) {
        return Ne_(Tensor::Full({}, scalar_value, dtype_, GetDevice()));
    }

    std::vector<Tensor> NonZeroNumpy() const;

    Tensor NonZero() const;

    bool All() const;

    bool Any() const;

    bool AllClose(const Tensor& other,
                  double rtol = 1e-5,
                  double atol = 1e-8) const;

    Tensor IsClose(const Tensor& other,
                   double rtol = 1e-5,
                   double atol = 1e-8) const;

    bool IsSame(const Tensor& other) const;

    template <typename T>
    std::vector<T> ToFlatVector() const {
        AssertTemplateDtype<T>();
        std::vector<T> values(NumElements());
        MemoryManager::MemcpyToHost(values.data(), Contiguous().GetDataPtr(),
                                    GetDevice(),
                                    GetDtype().ByteSize() * NumElements());
        return values;
    }

    inline bool IsContiguous() const {
        return DefaultStrides(shape_) == strides_;
    };

    Tensor Contiguous() const;

    Tensor Matmul(const Tensor& rhs) const;

    Tensor Solve(const Tensor& rhs) const;

    Tensor LeastSquares(const Tensor& rhs) const;

    Tensor Inverse() const;

    std::tuple<Tensor, Tensor, Tensor> SVD() const;

    inline SizeVector GetShape() const { return shape_; }

    inline const SizeVector& GetShapeRef() const { return shape_; }

    inline int64_t GetShape(int64_t dim) const {
        return shape_[shape_util::WrapDim(dim, NumDims())];
    }

    inline SizeVector GetStrides() const { return strides_; }

    inline const SizeVector& GetStridesRef() const { return strides_; }

    inline int64_t GetStride(int64_t dim) const {
        return strides_[shape_util::WrapDim(dim, NumDims())];
    }

    inline void* GetDataPtr() { return data_ptr_; }

    inline const void* GetDataPtr() const { return data_ptr_; }

    inline Dtype GetDtype() const { return dtype_; }

    Device GetDevice() const;

    inline std::shared_ptr<Blob> GetBlob() const { return blob_; }

    inline int64_t NumElements() const { return shape_.NumElements(); }

    inline int64_t NumDims() const { return shape_.size(); }

    template <typename T>
    void AssertTemplateDtype() const {
        if (!dtype_.IsObject() && Dtype::FromType<T>() != dtype_) {
            utility::LogError(
                    "Requested values have type {} but Tensor has type {}",
                    Dtype::FromType<T>().ToString(), dtype_.ToString());
        }
        if (dtype_.ByteSize() != sizeof(T)) {
            utility::LogError("Internal error: element size mismatch {} != {}",
                              dtype_.ByteSize(), sizeof(T));
        }
    }

    static SizeVector DefaultStrides(const SizeVector& shape);

    static std::pair<bool, SizeVector> ComputeNewStrides(
            const SizeVector& old_shape,
            const SizeVector& old_strides,
            const SizeVector& new_shape);

    DLManagedTensor* ToDLPack() const;

    static Tensor FromDLPack(const DLManagedTensor* dlmt);

    void AssertShape(const SizeVector& expected_shape) const;

protected:
    std::string ScalarPtrToString(const void* ptr) const;

protected:
    SizeVector shape_ = {0};

    SizeVector strides_ = {1};

    void* data_ptr_ = nullptr;

    Dtype dtype_ = Dtype::Undefined;

    std::shared_ptr<Blob> blob_ = nullptr;
};  // namespace core

template <>
inline Tensor::Tensor(const std::vector<bool>& init_vals,
                      const SizeVector& shape,
                      Dtype dtype,
                      const Device& device)
    : Tensor(shape, dtype, device) {
    // Check number of elements
    if (static_cast<int64_t>(init_vals.size()) != shape_.NumElements()) {
        utility::LogError(
                "Tensor initialization values' size {} does not match the "
                "shape {}",
                init_vals.size(), shape_.NumElements());
    }

    // Check data types
    AssertTemplateDtype<bool>();

    // std::vector<bool> possibly implements 1-bit-sized boolean storage.
    // Open3D uses 1-byte-sized boolean storage for easy indexing.
    std::vector<uint8_t> init_vals_uchar(init_vals.size());
    std::transform(init_vals.begin(), init_vals.end(), init_vals_uchar.begin(),
                   [](bool v) -> uint8_t { return static_cast<uint8_t>(v); });

    MemoryManager::MemcpyFromHost(blob_->GetDataPtr(), GetDevice(),
                                  init_vals_uchar.data(),
                                  init_vals_uchar.size() * dtype.ByteSize());
}

template <>
inline std::vector<bool> Tensor::ToFlatVector() const {
    AssertTemplateDtype<bool>();
    std::vector<bool> values(NumElements());
    std::vector<uint8_t> values_uchar(NumElements());
    MemoryManager::MemcpyToHost(values_uchar.data(), Contiguous().GetDataPtr(),
                                GetDevice(),
                                GetDtype().ByteSize() * NumElements());

    // std::vector<bool> possibly implements 1-bit-sized boolean storage.
    // Open3D uses 1-byte-sized boolean storage for easy indexing.
    std::transform(values_uchar.begin(), values_uchar.end(), values.begin(),
                   [](uint8_t v) -> bool { return static_cast<bool>(v); });
    return values;
}

template <>
inline bool Tensor::Item() const {
    if (shape_.size() != 0) {
        utility::LogError("Item only works for scalar Tensor of shape ()");
    }
    AssertTemplateDtype<bool>();
    uint8_t value;
    MemoryManager::MemcpyToHost(&value, data_ptr_, GetDevice(),
                                sizeof(uint8_t));
    return static_cast<bool>(value);
}

template <typename Scalar>
inline void Tensor::Fill(Scalar v) {
    DISPATCH_DTYPE_TO_TEMPLATE_WITH_BOOL(GetDtype(), [&]() {
        scalar_t casted_v = static_cast<scalar_t>(v);
        Tensor tmp(std::vector<scalar_t>({casted_v}), SizeVector({}),
                   GetDtype(), GetDevice());
        AsRvalue() = tmp;
    });
}

template <typename Object>
inline void Tensor::FillObject(const Object& v) {
    Tensor tmp(std::vector<Object>({v}), SizeVector({}), GetDtype(),
               GetDevice());
    AsRvalue() = tmp;
}

template <typename T>
inline Tensor operator+(T scalar_lhs, const Tensor& rhs) {
    return rhs + scalar_lhs;
}

template <typename T>
inline Tensor operator-(T scalar_lhs, const Tensor& rhs) {
    return Tensor::Full({}, scalar_lhs, rhs.GetDtype(), rhs.GetDevice()) - rhs;
}

template <typename T>
inline Tensor operator*(T scalar_lhs, const Tensor& rhs) {
    return rhs * scalar_lhs;
}

template <typename T>
inline Tensor operator/(T scalar_lhs, const Tensor& rhs) {
    return Tensor::Full({}, scalar_lhs, rhs.GetDtype(), rhs.GetDevice()) / rhs;
}

}  // namespace core
}  // namespace open3d