TensorKey.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <optional>
 
#include "open3d/utility/Logging.h"
 
namespace open3d {
namespace core {
 
// Avoids circular include.
class Tensor;
 
// Same as std::nullopt. Provides a similar Python slicing API.
constexpr std::nullopt_t None = std::nullopt;
 
class TensorKey {
public:
    static TensorKey Index(int64_t index);
 
    static TensorKey Slice(std::optional<int64_t> start,
                           std::optional<int64_t> stop,
                           std::optional<int64_t> step);
 
    static TensorKey IndexTensor(const Tensor& index_tensor);
 
    enum class TensorKeyMode { Index, Slice, IndexTensor };
    ~TensorKey() {}
 
    TensorKeyMode GetMode() const;
 
    std::string ToString() const;
 
public:
    int64_t GetIndex() const;
 
    int64_t GetStart() const;
 
    int64_t GetStop() const;
 
    int64_t GetStep() const;
 
    TensorKey InstantiateDimSize(int64_t dim_size) const;
 
    Tensor GetIndexTensor() const;
 
private:
    class Impl;
    class IndexImpl;
    class SliceImpl;
    class IndexTensorImpl;
    std::shared_ptr<Impl> impl_;
    TensorKey(const std::shared_ptr<Impl>& impl);
};
 
}  // namespace core
}  // namespace open3d