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

#include <algorithm>
#include <set>
#include <stdexcept>
#include <string>

#include "tenzr/common/error.h"

namespace tenzr::apiutils {

  int64_t CanonicalizeAxis(int64_t axis, size_t rank) {
    int64_t r = static_cast<int64_t>(rank);
    TENZR_CHECK(r > 0, "Cannot canonicalize axis for a 0-rank scalar");
    TENZR_CHECK(axis >= -r && axis < r, "Axis " + std::to_string(axis) + " out of bounds for tensor of rank " + std::to_string(rank));
    return (axis < 0) ? (axis + r) : axis;
  }

  std::vector<int64_t> CanonicalizeAxes(const std::vector<int64_t>& axes, size_t rank) {
    std::set<int64_t> unique_axes;
    for (int64_t ax : axes) {
      unique_axes.insert(CanonicalizeAxis(ax, rank));
    }
    return std::vector<int64_t>(unique_axes.begin(), unique_axes.end());
  }

  std::vector<int64_t> ComputeContiguousStrides(const Shape& shape) {
    size_t r = shape.Rank();
    if (r == 0) {
      return {};
    }
    std::vector<int64_t> strides(r, 1);
    for (int64_t i = static_cast<int64_t>(r) - 2; i >= 0; --i) {
      strides[static_cast<size_t>(i)] = strides[static_cast<size_t>(i + 1)] * shape.Dim(static_cast<size_t>(i + 1));
    }
    return strides;
  }

  int64_t ComputeElementOffset(const std::vector<int64_t>& indices, const std::vector<int64_t>& strides, int64_t base_offset) {
    TENZR_CHECK(indices.size() == strides.size(),
                "indices count (" + std::to_string(indices.size()) + ") must match strides count (" + std::to_string(strides.size()) + ")");
    int64_t off = base_offset;
    for (size_t i = 0; i < indices.size(); ++i) {
      off += indices[i] * strides[i];
    }
    return off;
  }

  bool AreBroadcastCompatible(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 < rank_a) ? a.Dim(rank_a - 1 - i) : 1;
      int64_t dim_b = (i < rank_b) ? b.Dim(rank_b - 1 - i) : 1;

      if (dim_a != dim_b && dim_a != 1 && dim_b != 1) {
        return false;
      }
    }
    return true;
  }

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

    std::vector<int64_t> out_dims(max_rank);
    for (size_t i = 0; i < max_rank; ++i) {
      int64_t dim_a = (i < rank_a) ? a.Dim(rank_a - 1 - i) : 1;
      int64_t dim_b = (i < rank_b) ? b.Dim(rank_b - 1 - i) : 1;

      if (dim_a != dim_b && dim_a != 1 && dim_b != 1) {
        TENZR_THROW("Incompatible shapes for broadcasting: " + a.ToString() + " and " + b.ToString());
      }
      out_dims[max_rank - 1 - i] = std::max(dim_a, dim_b);
    }
    return Shape(std::move(out_dims));
  }

  Shape BroadcastShapes(const std::vector<Shape>& shapes) {
    if (shapes.empty()) {
      return Shape({});
    }
    Shape cur = shapes[0];
    for (size_t i = 1; i < shapes.size(); ++i) {
      cur = BroadcastShapes(cur, shapes[i]);
    }
    return cur;
  }

  Shape SqueezeShape(const Shape& shape, const std::vector<int64_t>& axes) {
    std::vector<int64_t> canon = CanonicalizeAxes(axes, shape.Rank());
    std::set<int64_t> remove_set(canon.begin(), canon.end());

    std::vector<int64_t> new_dims;
    for (size_t i = 0; i < shape.Rank(); ++i) {
      if (remove_set.count(static_cast<int64_t>(i))) {
        TENZR_CHECK(shape.Dim(i) == 1, "Cannot squeeze axis " + std::to_string(i) + " with size " + std::to_string(shape.Dim(i)));
      } else {
        new_dims.push_back(shape.Dim(i));
      }
    }
    return Shape(std::move(new_dims));
  }

  Shape SqueezeAllUnitDims(const Shape& shape) {
    std::vector<int64_t> new_dims;
    for (int64_t d : shape.Dims()) {
      if (d != 1) {
        new_dims.push_back(d);
      }
    }
    return Shape(std::move(new_dims));
  }

  Shape ExpandDimsShape(const Shape& shape, int64_t axis) {
    size_t rank = shape.Rank();
    int64_t r = static_cast<int64_t>(rank);
    TENZR_CHECK(axis >= -r - 1 && axis <= r, "Axis " + std::to_string(axis) + " out of bounds for expand_dims with rank " + std::to_string(rank));
    int64_t canon_ax = (axis < 0) ? (axis + r + 1) : axis;

    size_t idx = static_cast<size_t>(canon_ax);
    std::vector<int64_t> new_dims;
    new_dims.reserve(rank + 1);
    for (size_t i = 0; i < rank; ++i) {
      if (i == idx) {
        new_dims.push_back(1);
      }
      new_dims.push_back(shape.Dim(i));
    }
    if (idx == rank) {
      new_dims.push_back(1);
    }
    return Shape(std::move(new_dims));
  }

  Shape FlattenShape(const Shape& shape) { return Shape({shape.NumElements()}); }

}  // namespace tenzr::apiutils