tensor_utils.h

/*
 * Copyright 2022 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
#pragma once
 
#include <lagrange/Logger.h>
#include <lagrange/utils/span.h>
 
#include <lagrange/Attribute.h>
#include <lagrange/IndexedAttribute.h>
#include <lagrange/python/binding.h>
#include <lagrange/utils/SmallVector.h>
 
#include <type_traits>
#include <vector>
 
namespace lagrange::python {
namespace nb = nanobind;
using namespace nb::literals;
 
template <typename ValueType>
using Tensor = nb::ndarray<ValueType, nb::numpy, nb::c_contig, nb::device::cpu>;
using GenericTensor = nb::ndarray<nb::c_contig, nb::numpy, nb::device::cpu>;
using Shape = SmallVector<size_t, 2>;
using Stride = SmallVector<int64_t, 2>;

bool is_vector(const Shape& shape);

bool check_shape(const Shape& shape, size_t expected_size);

bool check_shape(const Shape& shape, size_t expected_rows, size_t expected_cols);

bool is_dense(const Shape& shape, const Stride& stride);
 

template <typename ValueType>
Tensor<ValueType> create_empty_tensor();
 
template <typename ValueType>
std::tuple<span<ValueType>, Shape, Stride> tensor_to_span(Tensor<ValueType> tensor);
 
template <typename ValueType>
Tensor<std::decay_t<ValueType>> span_to_tensor(span<ValueType> values, nb::handle base);
 
template <typename ValueType>
Tensor<std::decay_t<ValueType>>
span_to_tensor(span<ValueType> values, span<const size_t> shape, nb::handle base);
 
template <typename ValueType>
Tensor<std::decay_t<ValueType>> span_to_tensor(
    span<ValueType> values,
    span<const size_t> shape,
    span<const int64_t> stride,
    nb::handle base);
 
template <typename ValueType>
Tensor<std::decay_t<ValueType>> attribute_to_tensor(
    const Attribute<ValueType>& attr,
    nb::handle base);
 
template <typename ValueType>
Tensor<std::decay_t<ValueType>>
attribute_to_tensor(const Attribute<ValueType>& attr, span<const size_t> shape, nb::handle base);
 
} // namespace lagrange::python