// 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.
#include "tenzr/frontend/layout.h"
#include <algorithm>
#include <cmath>
#include <numeric>
#include <sstream>
#include "tenzr/common/error.h"
namespace tenzr {
StridedLayout::StridedLayout(Shape shape) : shape_(shape), strides_(shape.ComputeContiguousStrides()), offset_(0) {}
StridedLayout::StridedLayout(Shape shape, std::vector<int64_t> strides, int64_t offset)
: shape_(std::move(shape)), strides_(std::move(strides)), offset_(offset) {
TENZR_CHECK(shape_.Rank() == strides_.size(), "Strides count must match shape rank");
}
int64_t StridedLayout::Stride(size_t index) const {
TENZR_CHECK(index < strides_.size(), "Stride index out of range: " + std::to_string(index) + " >= " + std::to_string(strides_.size()));
return strides_[index];
}
bool StridedLayout::IsCContiguous() const {
if (offset_ != 0) {
return false;
}
if (shape_.Rank() <= 1) {
return strides_.empty() || strides_.back() == 1;
}
auto expected = shape_.ComputeContiguousStrides();
return strides_ == expected;
}
bool StridedLayout::IsFContiguous() const {
if (offset_ != 0) {
return false;
}
if (shape_.Rank() <= 1) {
return strides_.empty() || strides_.front() == 1;
}
int64_t expected_stride = 1;
for (size_t i = 0; i < shape_.Rank(); ++i) {
if (strides_[i] != expected_stride) {
return false;
}
expected_stride *= shape_.Dim(i);
}
return true;
}
bool StridedLayout::IsContiguous() const { return IsCContiguous(); }
int64_t StridedLayout::LinearOffset(const std::vector<int64_t>& indices) const {
TENZR_CHECK(indices.size() == shape_.Rank(), "Indices count must match shape rank");
int64_t off = offset_;
for (size_t i = 0; i < indices.size(); ++i) {
int64_t idx = indices[i];
TENZR_CHECK(idx >= 0 && idx < shape_.Dim(i),
"Index out of bounds along axis " + std::to_string(i) + ": " + std::to_string(idx) + " not in [0, " + std::to_string(shape_.Dim(i)) + ")");
off += idx * strides_[i];
}
return off;
}
StridedLayout StridedLayout::Transpose(int64_t dim0, int64_t dim1) const {
const int64_t r = static_cast<int64_t>(shape_.Rank());
TENZR_CHECK(r >= 2, "Transpose requires rank >= 2, got rank " + std::to_string(r));
int64_t d0 = dim0 < 0 ? dim0 + r : dim0;
int64_t d1 = dim1 < 0 ? dim1 + r : dim1;
TENZR_CHECK(d0 >= 0 && d0 < r, "Invalid dim0: " + std::to_string(dim0));
TENZR_CHECK(d1 >= 0 && d1 < r, "Invalid dim1: " + std::to_string(dim1));
std::vector<int64_t> new_dims = shape_.Dims();
std::vector<int64_t> new_strides = strides_;
std::swap(new_dims[static_cast<size_t>(d0)], new_dims[static_cast<size_t>(d1)]);
std::swap(new_strides[static_cast<size_t>(d0)], new_strides[static_cast<size_t>(d1)]);
return StridedLayout(Shape(new_dims), std::move(new_strides), offset_);
}
StridedLayout StridedLayout::Permute(const std::vector<int64_t>& axes) const {
const size_t r = shape_.Rank();
TENZR_CHECK(axes.size() == r, "Permute axes size must match rank");
std::vector<bool> seen(r, false);
std::vector<int64_t> new_dims(r);
std::vector<int64_t> new_strides(r);
for (size_t i = 0; i < r; ++i) {
int64_t axis = axes[i];
if (axis < 0) {
axis += static_cast<int64_t>(r);
}
TENZR_CHECK(axis >= 0 && axis < static_cast<int64_t>(r), "Invalid permute axis: " + std::to_string(axes[i]));
TENZR_CHECK(!seen[static_cast<size_t>(axis)], "Duplicate permute axis: " + std::to_string(axis));
seen[static_cast<size_t>(axis)] = true;
new_dims[i] = shape_.Dim(static_cast<size_t>(axis));
new_strides[i] = strides_[static_cast<size_t>(axis)];
}
return StridedLayout(Shape(new_dims), std::move(new_strides), offset_);
}
StridedLayout StridedLayout::Squeeze(std::optional<int64_t> dim) const {
const int64_t r = static_cast<int64_t>(shape_.Rank());
std::vector<int64_t> new_dims;
std::vector<int64_t> new_strides;
if (dim.has_value()) {
int64_t d = dim.value() < 0 ? dim.value() + r : dim.value();
TENZR_CHECK(d >= 0 && d < r, "Squeeze dim out of range: " + std::to_string(dim.value()));
TENZR_CHECK(shape_.Dim(static_cast<size_t>(d)) == 1, "Cannot squeeze dimension with size " + std::to_string(shape_.Dim(static_cast<size_t>(d))));
for (int64_t i = 0; i < r; ++i) {
if (i != d) {
new_dims.push_back(shape_.Dim(static_cast<size_t>(i)));
new_strides.push_back(strides_[static_cast<size_t>(i)]);
}
}
} else {
for (size_t i = 0; i < shape_.Rank(); ++i) {
if (shape_.Dim(i) != 1) {
new_dims.push_back(shape_.Dim(i));
new_strides.push_back(strides_[i]);
}
}
}
return StridedLayout(Shape(new_dims), std::move(new_strides), offset_);
}
StridedLayout StridedLayout::Unsqueeze(int64_t dim) const {
const int64_t r = static_cast<int64_t>(shape_.Rank());
int64_t d = dim < 0 ? dim + r + 1 : dim;
TENZR_CHECK(d >= 0 && d <= r, "Unsqueeze dim out of range: " + std::to_string(dim));
std::vector<int64_t> new_dims;
std::vector<int64_t> new_strides;
new_dims.reserve(static_cast<size_t>(r + 1));
new_strides.reserve(static_cast<size_t>(r + 1));
int64_t next_stride = (d < r) ? (shape_.Dim(static_cast<size_t>(d)) * strides_[static_cast<size_t>(d)]) : 1;
for (int64_t i = 0; i <= r; ++i) {
if (i == d) {
new_dims.push_back(1);
new_strides.push_back(next_stride);
}
if (i < r) {
new_dims.push_back(shape_.Dim(static_cast<size_t>(i)));
new_strides.push_back(strides_[static_cast<size_t>(i)]);
}
}
return StridedLayout(Shape(new_dims), std::move(new_strides), offset_);
}
StridedLayout StridedLayout::slice(const std::vector<Slice>& slices) const {
TENZR_CHECK(slices.size() <= shape_.Rank(), "Too many slices provided: " + std::to_string(slices.size()) + " > " + std::to_string(shape_.Rank()));
std::vector<int64_t> new_dims = shape_.Dims();
std::vector<int64_t> new_strides = strides_;
int64_t new_offset = offset_;
for (size_t i = 0; i < slices.size(); ++i) {
const auto& s = slices[i];
int64_t d = shape_.Dim(i);
int64_t step = s.step;
TENZR_CHECK(step != 0, "Slice step cannot be 0");
int64_t start = s.start.value_or(step > 0 ? 0 : d - 1);
int64_t stop = s.stop.value_or(step > 0 ? d : -1);
if (start < 0) {
start += d;
}
if (stop < 0 && s.stop.has_value()) {
stop += d;
}
start = std::clamp(start, int64_t(0), d);
stop = std::clamp(stop, int64_t(-1), d);
int64_t len = 0;
if (step > 0 && start < stop) {
len = (stop - start + step - 1) / step;
} else if (step < 0 && start > stop) {
len = (start - stop + (-step) - 1) / (-step);
}
new_dims[i] = len;
new_offset += start * strides_[i];
new_strides[i] = strides_[i] * step;
}
return StridedLayout(Shape(new_dims), std::move(new_strides), new_offset);
}
StridedLayout StridedLayout::BroadcastTo(const Shape& target_shape) const {
const size_t src_rank = shape_.Rank();
const size_t tgt_rank = target_shape.Rank();
TENZR_CHECK(src_rank <= tgt_rank, "Cannot broadcast tensor to lower rank");
std::vector<int64_t> new_strides(tgt_rank, 0);
for (size_t i = 0; i < tgt_rank; ++i) {
size_t tgt_idx = tgt_rank - 1 - i;
int64_t tgt_dim = target_shape.Dim(tgt_idx);
if (i < src_rank) {
size_t src_idx = src_rank - 1 - i;
int64_t src_dim = shape_.Dim(src_idx);
if (src_dim == tgt_dim) {
new_strides[tgt_idx] = strides_[src_idx];
} else if (src_dim == 1) {
new_strides[tgt_idx] = 0; // Broadcast stride
} else {
TENZR_THROW("Incompatible broadcast dimension at axis " + std::to_string(src_idx) + ": " + std::to_string(src_dim) + " -> " +
std::to_string(tgt_dim));
}
} else {
new_strides[tgt_idx] = 0; // Prepending new outer dimension
}
}
return StridedLayout(target_shape, std::move(new_strides), offset_);
}
std::string StridedLayout::ToString() const {
std::ostringstream oss;
oss << "StridedLayout(shape=" << shape_.ToString() << ", strides=[";
for (size_t i = 0; i < strides_.size(); ++i) {
oss << strides_[i] << (i + 1 < strides_.size() ? ", " : "");
}
oss << "], offset=" << offset_ << ")";
return oss.str();
}
} // namespace tenzr