Tensor.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018-2021 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 <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/Scalar.h"
#include "open3d/core/ShapeUtil.h"
#include "open3d/core/SizeVector.h"
#include "open3d/core/TensorCheck.h"
#include "open3d/core/TensorInit.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_(shape_util::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) && {
        this->Fill(v);
        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 S>
    void Fill(S 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"));

    template <typename T>
    static Tensor Init(const T val, const Device& device = Device("CPU:0")) {
        Dtype type = Dtype::FromType<T>();
        std::vector<T> ele_list{val};
        SizeVector shape;
        return Tensor(ele_list, shape, type, device);
    }

    template <typename T>
    static Tensor Init(const std::initializer_list<T>& in_list,
                       const Device& device = Device("CPU:0")) {
        return InitWithInitializerList<T, 1>(in_list, device);
    }

    template <typename T>
    static Tensor Init(
            const std::initializer_list<std::initializer_list<T>>& in_list,
            const Device& device = Device("CPU:0")) {
        return InitWithInitializerList<T, 2>(in_list, device);
    }

    template <typename T>
    static Tensor Init(
            const std::initializer_list<
                    std::initializer_list<std::initializer_list<T>>>& in_list,
            const Device& device = Device("CPU:0")) {
        return InitWithInitializerList<T, 3>(in_list, device);
    }

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

    static Tensor Diag(const Tensor& input);

    static Tensor Arange(const Scalar start,
                         const Scalar stop,
                         const Scalar step = 1,
                         const Dtype dtype = core::Int64,
                         const Device& device = core::Device("CPU:0"));

    Tensor Reverse() const;

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

    Tensor Append(
            const Tensor& other,
            const utility::optional<int64_t>& axis = utility::nullopt) const;

    Tensor Broadcast(const SizeVector& dst_shape) const;

    Tensor Expand(const SizeVector& dst_shape) const;

    Tensor Reshape(const SizeVector& dst_shape) const;

    Tensor Flatten(int64_t start_dim = 0, int64_t end_dim = -1) const;

    Tensor View(const SizeVector& dst_shape) const;

    Tensor Clone() const { return To(GetDevice(), /*copy=*/true); }

    void CopyFrom(const Tensor& other);

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

    Tensor To(const Device& device, bool copy = false) const;

    Tensor To(const Device& device, 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() { return *this; }

    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;

    double Det() const;

    template <typename T>
    T Item() const {
        if (shape_.NumElements() != 1) {
            utility::LogError(
                    "Tensor::Item() only works for Tensor with exactly one "
                    "element.");
        }
        AssertTemplateDtype<T>();
        T value;
        MemoryManager::MemcpyToHost(&value, data_ptr_, GetDevice(), sizeof(T));
        return value;
    }

    Tensor Add(const Tensor& value) const;
    Tensor Add(Scalar value) const;
    Tensor operator+(const Tensor& value) const { return Add(value); }
    Tensor operator+(Scalar value) const { return Add(value); }

    Tensor Add_(const Tensor& value);
    Tensor Add_(Scalar value);
    Tensor operator+=(const Tensor& value) { return Add_(value); }
    Tensor operator+=(Scalar value) { return Add_(value); }

    Tensor Sub(const Tensor& value) const;
    Tensor Sub(Scalar value) const;
    Tensor operator-(const Tensor& value) const { return Sub(value); }
    Tensor operator-(Scalar value) const { return Sub(value); }

    Tensor Sub_(const Tensor& value);
    Tensor Sub_(Scalar value);
    Tensor operator-=(const Tensor& value) { return Sub_(value); }
    Tensor operator-=(Scalar value) { return Sub_(value); }

    Tensor Mul(const Tensor& value) const;
    Tensor Mul(Scalar value) const;
    Tensor operator*(const Tensor& value) const { return Mul(value); }
    Tensor operator*(Scalar value) const { return Mul(value); }

    Tensor Mul_(const Tensor& value);
    Tensor Mul_(Scalar value);
    Tensor operator*=(const Tensor& value) { return Mul_(value); }
    Tensor operator*=(Scalar value) { return Mul_(value); }

    Tensor Div(const Tensor& value) const;
    Tensor Div(Scalar value) const;
    Tensor operator/(const Tensor& value) const { return Div(value); }
    Tensor operator/(Scalar value) const { return Div(value); }

    Tensor Div_(const Tensor& value);
    Tensor Div_(Scalar value);
    Tensor operator/=(const Tensor& value) { return Div_(value); }
    Tensor operator/=(Scalar value) { return Div_(value); }

    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 IsNan() const;

    Tensor IsInf() const;

    Tensor IsFinite() const;

    Tensor Clip(Scalar min_val, Scalar max_val) const;

    Tensor Clip_(Scalar min_val, Scalar max_val);

    Tensor Floor() const;

    Tensor Ceil() const;

    Tensor Round() const;

    Tensor Trunc() const;

    Tensor LogicalNot() const;

    Tensor LogicalNot_();

    Tensor LogicalAnd(const Tensor& value) const;
    Tensor operator&&(const Tensor& value) const { return LogicalAnd(value); }
    Tensor LogicalAnd(Scalar value) const;

    Tensor LogicalAnd_(const Tensor& value);
    Tensor LogicalAnd_(Scalar value);

    Tensor LogicalOr(const Tensor& value) const;
    Tensor operator||(const Tensor& value) const { return LogicalOr(value); }
    Tensor LogicalOr(Scalar value) const;

    Tensor LogicalOr_(const Tensor& value);
    Tensor LogicalOr_(Scalar value);

    Tensor LogicalXor(const Tensor& value) const;
    Tensor LogicalXor(Scalar value) const;

    Tensor LogicalXor_(const Tensor& value);
    Tensor LogicalXor_(Scalar value);

    Tensor Gt(const Tensor& value) const;
    Tensor operator>(const Tensor& value) const { return Gt(value); }
    Tensor Gt(Scalar value) const;

    Tensor Gt_(const Tensor& value);
    Tensor Gt_(Scalar value);

    Tensor Lt(const Tensor& value) const;
    Tensor operator<(const Tensor& value) const { return Lt(value); }
    Tensor Lt(Scalar value) const;

    Tensor Lt_(const Tensor& value);
    Tensor Lt_(Scalar value);

    Tensor Ge(const Tensor& value) const;
    Tensor operator>=(const Tensor& value) const { return Ge(value); }
    Tensor Ge(Scalar value) const;

    Tensor Ge_(const Tensor& value);
    Tensor Ge_(Scalar value);

    Tensor Le(const Tensor& value) const;
    Tensor operator<=(const Tensor& value) const { return Le(value); }
    Tensor Le(Scalar value) const;

    Tensor Le_(const Tensor& value);
    Tensor Le_(Scalar value);

    Tensor Eq(const Tensor& value) const;
    Tensor operator==(const Tensor& value) const { return Eq(value); }
    Tensor Eq(Scalar value) const;

    Tensor Eq_(const Tensor& value);
    Tensor Eq_(Scalar value);

    Tensor Ne(const Tensor& value) const;
    Tensor operator!=(const Tensor& value) const { return Ne(value); }
    Tensor Ne(Scalar value) const;

    Tensor Ne_(const Tensor& value);
    Tensor Ne_(Scalar value);

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

    Tensor NonZero() const;

    bool IsNonZero() const;

    bool All() const;

    bool Any() const;

    bool AllEqual(const Tensor& other) 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 shape_util::DefaultStrides(shape_) == strides_;
    }

    Tensor Contiguous() const;

    Tensor Matmul(const Tensor& rhs) const;

    Tensor Solve(const Tensor& rhs) const;

    Tensor LeastSquares(const Tensor& rhs) const;

    std::tuple<Tensor, Tensor, Tensor> LU(const bool permute_l = false) const;

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

    Tensor Triu(const int diagonal = 0) const;

    Tensor Tril(const int diagonal = 0) const;

    std::tuple<Tensor, Tensor> Triul(const int diagonal = 0) const;

    Tensor Inverse() const;

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

    inline int64_t GetLength() const { return GetShape().GetLength(); }

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

    template <typename T>
    inline T* GetDataPtr() {
        return const_cast<T*>(const_cast<const Tensor*>(this)->GetDataPtr<T>());
    }

    template <typename T>
    inline const T* GetDataPtr() const {
        if (!dtype_.IsObject() && Dtype::FromType<T>() != dtype_) {
            utility::LogError(
                    "Requested values have type {} but Tensor has type {}. "
                    "Please use non templated GetDataPtr() with manual "
                    "casting.",
                    Dtype::FromType<T>().ToString(), dtype_.ToString());
        }
        return static_cast<T*>(data_ptr_);
    }

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

    DLManagedTensor* ToDLPack() const;

    static Tensor FromDLPack(const DLManagedTensor* dlmt);

    void Save(const std::string& file_name) const;

    static Tensor Load(const std::string& file_name);

    struct Iterator {
        using iterator_category = std::forward_iterator_tag;
        using difference_type = std::ptrdiff_t;
        using value_type = Tensor;
        using pointer = value_type*;
        using reference = value_type;  // Typically Tensor&, but a tensor slice
                                       // creates a new Tensor object with
                                       // shared memory.

        // Iterator must be constructible, copy-constructible, copy-assignable,
        // destructible and swappable.
        Iterator(pointer tensor, int64_t index);
        Iterator(const Iterator&);
        ~Iterator();
        reference operator*() const;
        pointer operator->() const;
        Iterator& operator++();
        Iterator operator++(int);
        bool operator==(const Iterator& other) const;
        bool operator!=(const Iterator& other) const;

    private:
        struct Impl;
        std::unique_ptr<Impl> impl_;
    };

    struct ConstIterator {
        using iterator_category = std::forward_iterator_tag;
        using difference_type = std::ptrdiff_t;
        using value_type = const Tensor;
        using pointer = value_type*;
        using reference = value_type;  // Typically Tensor&, but a tensor slice
                                       // creates a new Tensor object with
                                       // shared memory.

        // ConstIterator must be constructible, copy-constructible,
        // copy-assignable, destructible and swappable.
        ConstIterator(pointer tensor, int64_t index);
        ConstIterator(const ConstIterator&);
        ~ConstIterator();
        reference operator*() const;
        pointer operator->() const;
        ConstIterator& operator++();
        ConstIterator operator++(int);
        bool operator==(const ConstIterator& other) const;
        bool operator!=(const ConstIterator& other) const;

    private:
        struct Impl;
        std::unique_ptr<Impl> impl_;
    };

    Iterator begin();

    Iterator end();

    ConstIterator cbegin() const;

    ConstIterator cend() const;

    ConstIterator begin() const { return cbegin(); }

    ConstIterator end() const { return cend(); }

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

private:
    template <typename T, size_t D>
    static Tensor InitWithInitializerList(
            const tensor_init::NestedInitializerList<T, D>& nested_list,
            const Device& device = Device("CPU:0")) {
        SizeVector shape = tensor_init::InferShape(nested_list);
        std::vector<T> values =
                tensor_init::ToFlatVector<T, D>(shape, nested_list);
        return Tensor(values, shape, Dtype::FromType<T>(), device);
    }

protected:
    SizeVector shape_ = {0};

    SizeVector strides_ = {1};

    void* data_ptr_ = nullptr;

    Dtype dtype_ = core::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_.NumElements() != 1) {
        utility::LogError(
                "Tensor::Item only works for Tensor with one element.");
    }
    AssertTemplateDtype<bool>();
    uint8_t value;
    MemoryManager::MemcpyToHost(&value, data_ptr_, GetDevice(),
                                sizeof(uint8_t));
    return static_cast<bool>(value);
}

template <typename S>
inline void Tensor::Fill(S 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