include/tenzr/frontend/layout.h¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed 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
//
//      https://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 CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef TENZR_FRONTEND_LAYOUT_H_
#define TENZR_FRONTEND_LAYOUT_H_

#include <cstdint>
#include <optional>
#include <string>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/frontend/shape.h"

namespace tenzr {

  struct Slice {
    std::optional<int64_t> start{std::nullopt};
    std::optional<int64_t> stop{std::nullopt};
    int64_t step{1};

    static Slice All() { return Slice{std::nullopt, std::nullopt, 1}; }

    static Slice Index(int64_t idx) { return Slice{idx, idx + 1, 1}; }

    static Slice Range(int64_t start_idx, int64_t stop_idx, int64_t step_size = 1) { return Slice{start_idx, stop_idx, step_size}; }
  };

  class StridedLayout {
   public:
    StridedLayout() = default;

    explicit StridedLayout(Shape shape);

    StridedLayout(Shape shape, std::vector<int64_t> strides, int64_t offset = 0);

    const Shape& shape() const { return shape_; }

    const std::vector<int64_t>& Strides() const { return strides_; }

    int64_t Offset() const { return offset_; }

    size_t Rank() const { return shape_.Rank(); }

    int64_t Stride(size_t index) const;

    int64_t NumElements() const { return shape_.NumElements(); }

    bool IsContiguous() const;

    bool IsCContiguous() const;

    bool IsFContiguous() const;

    int64_t LinearOffset(const std::vector<int64_t>& indices) const;

    StridedLayout Transpose(int64_t dim0 = -1, int64_t dim1 = -2) const;

    StridedLayout Permute(const std::vector<int64_t>& axes) const;

    StridedLayout Squeeze(std::optional<int64_t> dim = std::nullopt) const;

    StridedLayout Unsqueeze(int64_t dim) const;

    StridedLayout slice(const std::vector<Slice>& slices) const;

    StridedLayout BroadcastTo(const Shape& target_shape) const;

    std::string ToString() const;

   private:
    Shape shape_;
    std::vector<int64_t> strides_;
    int64_t offset_{0};
  };

}  // namespace tenzr

#endif  // TENZR_FRONTEND_LAYOUT_H_