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 <cstddef>
#include <memory>
#include <string>

#include "Open3D/Core/Blob.h"
#include "Open3D/Core/DLPack/DLPackConverter.h"
#include "Open3D/Core/DLPack/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 {

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() * DtypeUtil::ByteSize(dtype), 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>();

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

    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)
        : Tensor(other.GetShape(),
                 other.GetStrides(),
                 const_cast<void*>(other.GetDataPtr()),
                 other.GetDtype(),
                 other.GetBlob()) {}

    Tensor(Tensor&& other)
        : Tensor(other.GetShape(),
                 other.GetStrides(),
                 other.GetDataPtr(),
                 other.GetDtype(),
                 other.GetBlob()) {}

    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 T>
    void Fill(T v);

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

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

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

    DLManagedTensor* ToDLPack() const { return dlpack::ToDLPack(*this); }

    static Tensor FromDLPack(DLManagedTensor* src) {
        return dlpack::FromDLPack(src);
    }

    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;

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

    Tensor LogicalAnd_(const Tensor& value);

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

    Tensor LogicalOr_(const Tensor& value);

    Tensor LogicalXor(const Tensor& value) const;

    Tensor LogicalXor_(const Tensor& value);

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

    Tensor Gt_(const Tensor& value);

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

    Tensor Lt_(const Tensor& value);

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

    Tensor Ge_(const Tensor& value);

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

    Tensor Le_(const Tensor& value);

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

    Tensor Eq_(const Tensor& value);

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

    Tensor Ne_(const Tensor& value);

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

    Tensor NonZero() const;

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

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

    Tensor Contiguous() 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 (DtypeUtil::FromType<T>() != dtype_) {
            utility::LogError(
                    "Requested values have type {} but Tensor has type {}",
                    DtypeUtil::ToString(DtypeUtil::FromType<T>()),
                    DtypeUtil::ToString(dtype_));
        }
        if (DtypeUtil::ByteSize(dtype_) != sizeof(T)) {
            utility::LogError("Internal error: element size mismatch {} != {}",
                              DtypeUtil::ByteSize(dtype_), 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);

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 open3d

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<unsigned char> init_vals_uchar(init_vals.size());
    std::transform(init_vals.begin(), init_vals.end(), init_vals_uchar.begin(),
                   [](bool v) -> unsigned char {
                       return static_cast<unsigned char>(v);
                   });

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

template <>
inline std::vector<bool> Tensor::ToFlatVector() const {
    AssertTemplateDtype<bool>();
    std::vector<bool> values(NumElements());
    std::vector<unsigned char> values_uchar(NumElements());
    MemoryManager::MemcpyToHost(
            values_uchar.data(), Contiguous().GetDataPtr(), GetDevice(),
            DtypeUtil::ByteSize(GetDtype()) * 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(),
            [](unsigned char v) -> bool { return static_cast<bool>(v); });
    return values;
}

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 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 open3d