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 "open3d/utility/Logging.h"
#include "open3d/utility/Optional.h"
 
namespace open3d {
namespace core {
 
// Avoids circular include.
class Tensor;
 
// Same as utility::nullopt. Provides a similar Python slicing API.
constexpr utility::nullopt_t None{utility::nullopt_t::init()};

class TensorKey {
public:
    static TensorKey Index(int64_t index);

    static TensorKey Slice(utility::optional<int64_t> start,
                           utility::optional<int64_t> stop,
                           utility::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