src/frontend/shape.cc¶

// 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/shape.h"

#include <numeric>
#include <sstream>

#include "tenzr/common/error.h"

namespace tenzr {

  Shape::Shape(std::initializer_list<int64_t> dims) : dims_(dims) {}

  Shape::Shape(std::vector<int64_t> dims) : dims_(std::move(dims)) {}

  int64_t Shape::Dim(size_t index) const {
    TENZR_CHECK(index < dims_.size(), "Shape index out of range: " + std::to_string(index) + " >= " + std::to_string(dims_.size()));
    return dims_[index];
  }

  int64_t Shape::NumElements() const {
    if (dims_.empty()) {
      return 1;  // Scalar
    }
    int64_t total = 1;
    for (int64_t d : dims_) {
      TENZR_CHECK(d >= 0, "Negative dimension not allowed: " + std::to_string(d));
      total *= d;
    }
    return total;
  }

  std::vector<int64_t> Shape::ComputeContiguousStrides() const {
    if (dims_.empty()) {
      return {};
    }
    std::vector<int64_t> strides(dims_.size(), 1);
    for (size_t i = dims_.size() - 1; i > 0; --i) {
      strides[i - 1] = strides[i] * dims_[i];
    }
    return strides;
  }

  std::string Shape::ToString() const {
    std::ostringstream oss;
    oss << "[";
    for (size_t i = 0; i < dims_.size(); ++i) {
      if (i > 0) {
        oss << ", ";
      }
      oss << dims_[i];
    }
    oss << "]";
    return oss.str();
  }

  std::string Shape::ToMlirTensorType(DType dtype) const {
    std::ostringstream oss;
    oss << "tensor<";
    for (size_t i = 0; i < dims_.size(); ++i) {
      oss << dims_[i] << "x";
    }
    oss << ToMlirType(dtype) << ">";
    return oss.str();
  }

  Shape Shape::Broadcast(const Shape& a, const Shape& b) {
    size_t rank_a = a.Rank();
    size_t rank_b = b.Rank();
    size_t out_rank = std::max(rank_a, rank_b);
    std::vector<int64_t> out_dims(out_rank);

    for (size_t i = 0; i < out_rank; ++i) {
      int64_t dim_a = (i < out_rank - rank_a) ? 1 : a.Dim(i - (out_rank - rank_a));
      int64_t dim_b = (i < out_rank - rank_b) ? 1 : b.Dim(i - (out_rank - rank_b));
      if (dim_a == dim_b) {
        out_dims[i] = dim_a;
      } else if (dim_a == 1) {
        out_dims[i] = dim_b;
      } else if (dim_b == 1) {
        out_dims[i] = dim_a;
      } else {
        TENZR_THROW("Incompatible shapes for broadcasting: " + a.ToString() + " and " + b.ToString());
      }
    }
    return Shape(out_dims);
  }

  bool Shape::CanBroadcast(const Shape& a, const Shape& b) noexcept {
    size_t rank_a = a.Rank();
    size_t rank_b = b.Rank();
    size_t max_rank = std::max(rank_a, rank_b);

    for (size_t i = 0; i < max_rank; ++i) {
      int64_t dim_a = (i < max_rank - rank_a) ? 1 : a.Dim(i - (max_rank - rank_a));
      int64_t dim_b = (i < max_rank - rank_b) ? 1 : b.Dim(i - (max_rank - rank_b));
      if (dim_a != dim_b && dim_a != 1 && dim_b != 1) {
        return false;
      }
    }
    return true;
  }

}  // namespace tenzr