/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
*  See the NOTICE file distributed with this work for additional
*  information regarding copyright ownership.
* 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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/

//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <flatbuffers/util.h>
#include <helpers/ShapeUtils.h>

#include <algorithm>
#include <climits>
#include <numeric>
#include <set>

namespace sd {

//////////////////////////////////////////////////////////////////////////

LongType* ShapeUtils::evalTransposeShapeInfo(NDArray& arr, memory::Workspace* workspace,
                                             const bool setContigStrides) {
  LongType rank = arr.rankOf();

  // note we do this because of stack allocation crashes
  //  if the stack is used a vector's data can cause crashes when it goes out of scope
  LongType* dims = new LongType[rank];
  for (LongType i = 0; i < rank; i++) {
    dims[i] = rank - 1 - i;
  }

  auto ret = evalPermShapeInfo(dims, rank, &arr, workspace, setContigStrides);
  delete[] dims;
  return ret;
}

// evaluate shape for array resulting from tensorDot operation, also evaluate shapes and dimensions permutations for
// transposition of two input arrays
std::vector<LongType> ShapeUtils::evalShapeForTensorDot( LongType* aShapeInfo,  LongType* bShapeInfo,
                                                         const std::vector<LongType> axesA,
                                                         const std::vector<LongType> axesB,
                                                         std::vector<LongType>& permutAt,
                                                         std::vector<LongType>& permutBt, std::vector<LongType>& shapeAt,
                                                         std::vector<LongType>& shapeBt) {
  LongType axeAsize = static_cast<LongType>(axesA.size());
  LongType axeBsize = static_cast<LongType>(axesB.size());

  LongType aRank = aShapeInfo[0];
  LongType bRank = bShapeInfo[0];

  if (axeAsize != axeBsize) {
    std::string errorMessage;
    errorMessage +=
        "ShapeUtils::evalShapeForTensorDot method: the numbers of a axes and b axes to make dot product along must "
        "have identical values !\n";
    errorMessage += "axesASize: ";
    errorMessage += std::to_string(axeAsize);
    errorMessage += ", axesBSize: ";
    errorMessage += std::to_string(axeBsize);
    errorMessage += "\n";
    THROW_EXCEPTION(errorMessage.c_str());
  }

  if (axeAsize > aRank || axeBsize > bRank) {
    std::string errorMessage;
    errorMessage +=
        "ShapeUtils::evalShapeForTensorDot method: the length of vector of a or b axes is larger than array rank !\n";
    errorMessage += "axesASize: ";
    errorMessage += std::to_string(axeAsize);
    errorMessage += ", axesBSize: ";
    errorMessage += std::to_string(axeBsize);
    errorMessage += "\n";
    errorMessage += "aRank: ";
    errorMessage += std::to_string(aRank);
    errorMessage += ", bRank: ";
    errorMessage += std::to_string(bRank);
    errorMessage += "\n";
    THROW_EXCEPTION(errorMessage.c_str());
  }

  // check whether axesA and axesB contain only unique numbers
  std::set<LongType> uniqueElems(axesA.begin(), axesA.end());
  if ((LongType)uniqueElems.size() != axeAsize) {
    THROW_EXCEPTION("ShapeUtils::evalShapeForTensorDot method: the vector of a axes contains duplicates !");
  }
  uniqueElems.clear();
  uniqueElems = std::set<LongType>(axesB.begin(), axesB.end());
  if ((LongType)uniqueElems.size() != axeBsize) {
    std::string errorMessage;
    errorMessage += "ShapeUtils::evalShapeForTensorDot method: the vector of b axes contains duplicates !\n";
    errorMessage += "axesBsize: ";
    errorMessage += std::to_string(axesB.size());
    errorMessage += " uniqueElems: ";
    errorMessage += std::to_string(uniqueElems.size());
    THROW_EXCEPTION(errorMessage.c_str());
  }
  std::vector<LongType> list_A, list_B;
  for (LongType i = 0; i < aRank; i++)
    if (std::find(axesA.begin(), axesA.end(), i) == axesA.end()) list_A.emplace_back(i);
  for (LongType i = 0; i < bRank; i++)
    if (std::find(axesB.begin(), axesB.end(), i) == axesB.end()) list_B.emplace_back(i);

  permutAt = list_A;
  permutAt.insert(permutAt.end(), axesA.begin(), axesA.end());
  permutBt = axesB;
  permutBt.insert(permutBt.end(), list_B.begin(), list_B.end());

  // if permute contains something like {0,1,2,..rank-1}, then there is no need to make permutation and we return empty
  // vector in this case
  LongType i1, i2;
  for (i1 = 0; i1 < aRank; ++i1)
    if (permutAt[i1] != i1) break;
  if (i1 == aRank) permutAt = {};
  for (i2 = 0; i2 < bRank; ++i2)
    if (permutBt[i2] != i2) break;
  if (i2 == bRank) permutBt = {};

  LongType n2 = 1;
  for (LongType i = 0; i < axeAsize; i++) n2 *= aShapeInfo[axesA[i] + 1];
  shapeAt = {shape::length(aShapeInfo) / n2, n2};

  std::vector<LongType> oldShapeA;
  oldShapeA.resize(list_A.size());
  for (size_t i = 0; i < oldShapeA.size(); ++i) oldShapeA[i] = aShapeInfo[list_A[i] + 1];

  LongType n3 = 1;
  for (LongType i = 0; i < axeBsize; i++) n3 *= bShapeInfo[axesB[i] + 1];
  shapeBt = {n3, shape::length(bShapeInfo) / n3};

  std::vector<LongType> oldShapeB;
  oldShapeB.resize(list_B.size());
  for (size_t i = 0; i < oldShapeB.size(); i++) oldShapeB[i] = bShapeInfo[list_B[i] + 1];

  std::vector<LongType> aPlusB(oldShapeA);
  aPlusB.insert(aPlusB.end(), oldShapeB.begin(), oldShapeB.end());

  return aPlusB;
}

//////////////////////////////////////////////////////////////////////////
std::vector<LongType> ShapeUtils::evalShapeForTensorDot(NDArray* a, NDArray* b,
                                                        const std::vector<LongType>& axesA,
                                                        const std::vector<LongType>& axesB,
                                                        std::vector<LongType>& permutAt,
                                                        std::vector<LongType>& permutBt, std::vector<LongType>& shapeAt,
                                                        std::vector<LongType>& shapeBt) {
  return evalShapeForTensorDot(a->shapeInfo(), b->shapeInfo(), axesA, axesB, permutAt, permutBt, shapeAt, shapeBt);
}

//////////////////////////////////////////////////////////////////////////
// evaluate output shape for reduce operation when input shape is empty
LongType* ShapeUtils::evalReduceShapeInfoEmpty(const char order, std::vector<LongType>* dimsToExclude,
                                               LongType* shapeInfo, const DataType dataType,
                                               const bool keepDims, memory::Workspace* workspace) {
  if (dimsToExclude->size() == 0) {  // return copy of input shape
    LongType* outShapeInfo = ShapeBuilders::copyShapeInfoAndType(shapeInfo, dataType, true, workspace);
    auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(outShapeInfo)->primary();
    return ret;
  }

  const LongType rank = shape::rank(shapeInfo);
  LongType* outShapeInfo = nullptr;

  if (static_cast<sd::LongType>(dimsToExclude->size()) == rank) {  // return scalar or shape filled with unities

    if (!keepDims)
      outShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
    else
      outShapeInfo = ShapeBuilders::createShapeInfo(dataType, order, std::vector<LongType>(rank, 1), workspace);
  } else {
    shape::checkDimensions(rank, dimsToExclude);

    std::vector<LongType> outShape;

    if (keepDims) {
      outShape.assign(shapeInfo + 1, shapeInfo + 1 + rank);
      for (const auto dim : *dimsToExclude) outShape[dim] = 1;
    } else {
      for (LongType i = 0, j = 0; i < rank; ++i) {
        if (j < static_cast<sd::LongType>(dimsToExclude->size()) && i == dimsToExclude->at(j))
          ++j;
        else
          outShape.emplace_back(shapeInfo[i + 1]);
      }
    }

    outShapeInfo = ShapeBuilders::createShapeInfo(dataType, order, outShape, workspace);
  }


  auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(outShapeInfo)->primary();
  delete[] outShapeInfo;
  return ret;
}

LongType* ShapeUtils::evalReduceShapeInfo(const char order, std::vector<LongType>* dimsToExclude,
                                          NDArray& arr, const bool keepDims, const bool supportOldShapes,
                                          memory::Workspace* workspace) {
  return const_cast<LongType*>(
      evalReduceShapeInfo(order, dimsToExclude, arr, arr.dataType(), keepDims, supportOldShapes, workspace));
}

//////////////////////////////////////////////////////////////////////////
// return new (shorter) sorted dimensions array without dimensions that are present in input vector
std::vector<LongType>* ShapeUtils::evalDimsToExclude(const LongType rank, const LongType dimsLen, const LongType* dimensions) {
  std::vector<LongType> * ret = new std::vector<LongType>();



  // Validate input parameters
  if (rank <= 0) {
    THROW_EXCEPTION("ShapeUtils::evalDimsToExclude: rank must be positive");
  }
  if (dimsLen < 0) {
    THROW_EXCEPTION("ShapeUtils::evalDimsToExclude: dimsLen cannot be negative");
  }
  if (dimsLen > 0 && dimensions == nullptr) {
    THROW_EXCEPTION("ShapeUtils::evalDimsToExclude: dimensions array is null but dimsLen > 0");
  }

  if (dimsLen == 0) {  // if input vector is empty then return whole shape range
    ret->resize(rank);
    std::iota(ret->begin(), ret->end(), 0);  // fill with 0, 1, ... rank-1
  } else {
    // Validate dimensions are within bounds
    for (LongType j = 0; j < dimsLen; j++) {
      LongType dim = dimensions[j] >= 0 ? dimensions[j] : dimensions[j] + rank;
      if (dim < 0 || dim >= rank) {
        delete ret;
        THROW_EXCEPTION("ShapeUtils::evalDimsToExclude: dimension index is out of bounds");
      }
    }

    bool isAbsent;
    for (LongType i = 0; i < rank; i++) {
      isAbsent = true;
      for (LongType j = 0; j < dimsLen; j++) {
        LongType dim = dimensions[j] >= 0 ? dimensions[j] : dimensions[j] + rank;
        if (i == dim) {
          isAbsent = false;
          break;
        }
      }
      if (isAbsent) ret->emplace_back(i);
    }
  }

  // Note: We keep the original behavior - if ret is empty, it means all dimensions
  // were excluded, which is a valid case that the caller (gather operation) should handle
  return ret;
}

//////////////////////////////////////////////////////////////////////////
// evaluate shape resulting from reduce operation
LongType* ShapeUtils::evalReduceShapeInfo(const char order, std::vector<LongType>* dimsToExclude, LongType* shapeInfo,
                                          const DataType dataType,
                                          const bool keepDims,
                                          const bool supportOldShapes,
                                          memory::Workspace* workspace) {
  if (ArrayOptions::arrayType(shapeInfo) == EMPTY) {
    return evalReduceShapeInfoEmpty(order, dimsToExclude, shapeInfo, dataType, keepDims, workspace);
  }
  LongType* newShapeInfo = nullptr;

  LongType rank = shape::rank(const_cast<LongType*>(shapeInfo));

  if (dimsToExclude->size() == 0) {  // return scalar or array with len=1 in this case
    if (keepDims && rank > 1) {
      newShapeInfo = new LongType[shape::shapeInfoLength(rank)];
      newShapeInfo[0] = rank;
      for (LongType i = 0; i < rank; ++i) newShapeInfo[i + 1] = 1;
      updateStridesAndType(newShapeInfo, shapeInfo, order);
      ArrayOptions::setDataType(newShapeInfo, dataType);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else if (supportOldShapes) {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    }
  }

  shape::checkDimensions(rank, dimsToExclude);

  LongType dimSize = dimsToExclude->size();

  if (keepDims) {
    newShapeInfo = new LongType[shape::shapeInfoLength(rank)];
    newShapeInfo[0] = rank;

    for (LongType i = 0; i < rank; ++i) {
      if (std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                             i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
        newShapeInfo[i + 1] = 1;
      else
        newShapeInfo[i + 1] = shapeInfo[i + 1];
    }
    updateStridesAndType(newShapeInfo, shapeInfo, order);
    auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
    delete[] newShapeInfo;
    return ret;
  }

  LongType newRank = rank - dimSize;
  if (newRank == 0 ||
      (dimSize == 1 &&
       dimsToExclude->at(0) == INT_MAX)) {  // check whether given dimension is meant for the whole dimension

    if (supportOldShapes) {
      newShapeInfo = new LongType[shape::shapeInfoLength(2)];
      shape::shapeOldScalar(ArrayOptions::dataType(shapeInfo), newShapeInfo, 'c');
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else {

      newShapeInfo = ShapeBuilders::createScalarShapeInfo(ArrayOptions::dataType(shapeInfo), workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    }
  }

  newShapeInfo = new LongType[shape::shapeInfoLength(newRank)];
  newShapeInfo[0] = newRank;  // set rank
  LongType j = 1;
  for (LongType i = 0; i < rank; ++i)
    if (!std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                            i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
      newShapeInfo[j++] = shapeInfo[i + 1];

  // ensure whether vector has proper shape for old shape type
  if (newRank == 1 && supportOldShapes) {
    LongType oldValue = newShapeInfo[1];
    delete[] newShapeInfo;
    newShapeInfo = new LongType[shape::shapeInfoLength(2)];
    newShapeInfo[0] = 2;
    if (dimsToExclude->at(0) == 0) {
      newShapeInfo[1] = 1;
      newShapeInfo[2] = oldValue;
    } else {
      newShapeInfo[1] = oldValue;
      newShapeInfo[2] = 1;
    }
  }

  updateStridesAndType(newShapeInfo, shapeInfo, order);

  auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
  delete[] newShapeInfo;
  return ret;
}


LongType* ShapeUtils::evalReduceShapeInfo(const char order, std::vector<LongType>* dimsToExclude, NDArray& arr,
                                          const DataType dataType, const bool keepDims,
                                          const bool supportOldShapes, memory::Workspace* workspace) {
  sd::LongType  *shapeInfo = arr.shapeInfo();
  if (ArrayOptions::arrayType(shapeInfo) == EMPTY)
    return evalReduceShapeInfoEmpty(order, dimsToExclude, shapeInfo, dataType, keepDims, workspace);

  LongType* newShapeInfo = nullptr;

  LongType rank = shape::rank(const_cast<LongType*>(shapeInfo));

  if (dimsToExclude->size() == 0) {  // return scalar or array with len=1 in this case
    if (keepDims && rank > 1) {
      newShapeInfo = new sd::LongType[shape::shapeInfoLength(rank)];
      newShapeInfo[0] = rank;
      for (LongType i = 0; i < rank; ++i) newShapeInfo[i + 1] = 1;
      updateStridesAndType(newShapeInfo, shapeInfo, order);
      ArrayOptions::setDataType(newShapeInfo, dataType);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else if (supportOldShapes) {
      newShapeInfo = new sd::LongType[shape::shapeInfoLength(2)];
      shape::shapeOldScalar(dataType, newShapeInfo, 'c');
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    }
  }

  shape::checkDimensions(rank, dimsToExclude);
  LongType dimSize = dimsToExclude->size();

  if (keepDims) {
    newShapeInfo = new sd::LongType[shape::shapeInfoLength(rank)];
    newShapeInfo[0] = rank;

    for (LongType i = 0; i < rank; ++i) {
      if (std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                             i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
        newShapeInfo[i + 1] = 1;
      else
        newShapeInfo[i + 1] = shapeInfo[i + 1];
    }
    updateStridesAndType(newShapeInfo, shapeInfo, order);
    auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
    delete[] newShapeInfo;
    return ret;
  }

  LongType newRank = rank - dimSize;
  if (newRank == 0 ||
      (dimSize == 1 &&
       dimsToExclude->at(0) == INT_MAX)) {  // check whether given dimension is meant for the whole dimension

    if (supportOldShapes) {
      newShapeInfo = new sd::LongType[shape::shapeInfoLength(2)];
      shape::shapeOldScalar(ArrayOptions::dataType(shapeInfo), newShapeInfo, 'c');
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;

      return ret;
    } else {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(ArrayOptions::dataType(shapeInfo), workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;

      return ret;
    }
  }

  newShapeInfo = new sd::LongType[shape::shapeInfoLength(newRank)];
  newShapeInfo[0] = newRank;  // set rank
  LongType j = 1;
  for (LongType i = 0; i < rank; ++i)
    if (!std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                            i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
      newShapeInfo[j++] = shapeInfo[i + 1];

  // ensure whether vector has proper shape for old shape type
  if (newRank == 1 && supportOldShapes) {
    LongType oldValue = newShapeInfo[1];
    delete[] newShapeInfo;
    RELEASE(newShapeInfo, workspace);
    newShapeInfo = new sd::LongType[shape::shapeInfoLength(2)];
    ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(2), sd::LongType);  // set newRank = 2
    newShapeInfo[0] = 2;
    if (dimsToExclude->at(0) == 0) {
      newShapeInfo[1] = 1;
      newShapeInfo[2] = oldValue;
    } else {
      newShapeInfo[1] = oldValue;
      newShapeInfo[2] = 1;
    }
  }

  updateStridesAndType(newShapeInfo, shapeInfo, order);

  auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
  RELEASE(newShapeInfo, workspace);
  return ret;


}
LongType* ShapeUtils::evalReduceShapeInfo(char order, std::vector<LongType>* dimsToExclude, LongType* shapeInfo, const bool keepDims,
                                          bool supportOldShapes, memory::Workspace* workspace) {


  sd::DataType dataType = ArrayOptions::dataType(shapeInfo);
  if (ArrayOptions::arrayType(shapeInfo) == EMPTY)
    return evalReduceShapeInfoEmpty(order, dimsToExclude, shapeInfo, dataType, keepDims, workspace);

  LongType* newShapeInfo = nullptr;

  LongType rank = shape::rank(const_cast<LongType*>(shapeInfo));

  if (dimsToExclude->size() == 0) {  // return scalar or array with len=1 in this case
    if (keepDims && rank > 1) {
      newShapeInfo = new sd::LongType[shape::shapeInfoLength(rank)];
      newShapeInfo[0] = rank;
      for (LongType i = 0; i < rank; ++i) newShapeInfo[i + 1] = 1;
      updateStridesAndType(newShapeInfo, shapeInfo, order);
      ArrayOptions::setDataType(newShapeInfo, dataType);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else if (supportOldShapes) {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
      shape::shapeOldScalar(dataType, newShapeInfo, 'c');
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else {
      newShapeInfo = ShapeBuilders::createScalarShapeInfo(dataType, workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    }
  }

  shape::checkDimensions(rank, dimsToExclude);

  LongType dimSize = dimsToExclude->size();

  if (keepDims) {
    newShapeInfo = new sd::LongType[shape::shapeInfoLength(rank)];
    newShapeInfo[0] = rank;

    for (LongType i = 0; i < rank; ++i) {
      if (std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                             i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
        newShapeInfo[i + 1] = 1;
      else
        newShapeInfo[i + 1] = shapeInfo[i + 1];
    }
    updateStridesAndType(newShapeInfo, shapeInfo, order);
    auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
    delete[] newShapeInfo;
    return ret;
  }

  LongType newRank = rank - dimSize;
  if (newRank == 0 ||
      (dimSize == 1 &&
       dimsToExclude->at(0) == INT_MAX)) {  // check whether given dimension is meant for the whole dimension

    if (supportOldShapes) {
      newShapeInfo = new sd::LongType[shape::shapeInfoLength(2)];
      shape::shapeOldScalar(ArrayOptions::dataType(shapeInfo), newShapeInfo, 'c');
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    } else {

      newShapeInfo = ShapeBuilders::createScalarShapeInfo(ArrayOptions::dataType(shapeInfo), workspace);
      auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
      delete[] newShapeInfo;
      return ret;
    }
  }

  newShapeInfo = new sd::LongType[shape::shapeInfoLength(newRank)];
  newShapeInfo[0] = newRank;  // set rank
  LongType j = 1;
  for (LongType i = 0; i < rank; ++i)
    if (!std::binary_search(dimsToExclude->begin(), dimsToExclude->end(),
                            i))  // dimsToExclude is already sorted after shape::checkDimensions() has been applied
      newShapeInfo[j++] = shapeInfo[i + 1];

  // ensure whether vector has proper shape for old shape type
  if (newRank == 1 && supportOldShapes) {
    LongType oldValue = newShapeInfo[1];
    delete[] newShapeInfo;
    newShapeInfo = new sd::LongType[shape::shapeInfoLength(2)];
    newShapeInfo[0] = 2;
    if (dimsToExclude->at(0) == 0) {
      newShapeInfo[1] = 1;
      newShapeInfo[2] = oldValue;
    } else {
      newShapeInfo[1] = oldValue;
      newShapeInfo[2] = 1;
    }
  }

  updateStridesAndType(newShapeInfo, shapeInfo, order);

  auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
  delete[] newShapeInfo;
  return ret;
}


//////////////////////////////////////////////////////////////////////////
// evaluate shape for array which is result of repeat operation applied to arr
std::vector<LongType> ShapeUtils::evalRepeatShape(LongType axis, const std::vector<LongType>& repeats,
                                                  NDArray& arr) {
  if (axis < 0) axis += arr.rankOf();

  if (repeats.size() != 1 && static_cast<LongType>(repeats.size()) != arr.sizeAt(axis))
    THROW_EXCEPTION(
        "ShapeUtils::evalRepeatShape: size of repeats vector must be 1 or equal to dimension at given axis !");

  auto* shapeVec = arr.getShapeAsVector();
  std::vector<LongType> outShape = *shapeVec;
  delete shapeVec;

  if (repeats.size() == 1)
    outShape[axis] *= repeats[0];
  else
    outShape[axis] = std::accumulate(repeats.begin(), repeats.end(), 0);

  return outShape;
}

//////////////////////////////////////////////////////////////////////////
// evaluate shapeInfo of permuted array
LongType* ShapeUtils::evalPermShapeInfo(LongType* dimensions, LongType rank, NDArray* arr,
                                        memory::Workspace* workspace, const bool setContigStrides) {
  if (rank != arr->rankOf())
    THROW_EXCEPTION("ShapeUtils::evalPermShapeInfo static method: wrong arguments: rank is not suitable!");

  auto shapeInfoLength = shape::shapeInfoLength(rank);

  // allocate memory for new array - shapeInfo
  LongType* shapeInfoNew = nullptr;
  ALLOCATE(shapeInfoNew, workspace, shapeInfoLength, sd::LongType);

  // copy arr _shapeInfo into new array
  memcpy(shapeInfoNew, arr->shapeInfo(), shape::shapeInfoByteLength(rank));

  // perform buffer permutation
  shape::doPermuteShapeInfo(shapeInfoNew, dimensions, rank);

  if (setContigStrides) {
    shape::updateStrides(shapeInfoNew, arr->ordering(), true);
  }


  shape::setOrder(shapeInfoNew, arr->ordering());
  ArrayOptions::setDataType(shapeInfoNew, arr->dataType());
  return shapeInfoNew;
}

//////////////////////////////////////////////////////////////////////////
bool ShapeUtils::copyVectorPart(std::vector<LongType>& target, std::vector<LongType>& source, LongType rank,
                                LongType offset) {
  if (static_cast<sd::LongType>(source.size()) < offset + rank) return false;

  for (LongType e = offset; e < offset + rank; e++) target.push_back(source[e]);

  return true;
}

//////////////////////////////////////////////////////////////////////////
// check whether 2 arrays have mutually broadcastable shapes
// shape comparison starts from the end
bool ShapeUtils::areShapesBroadcastable(NDArray& arr1, NDArray& arr2) {
  return areShapesBroadcastable(arr1.shapeInfo(), arr2.shapeInfo());
}

bool ShapeUtils::areShapesBroadcastable(const LongType* shapeInfo1, const LongType* shapeInfo2) {
  // Scalars can be broadcast with anything
  if (shape::isScalar(shapeInfo1) || shape::isScalar(shapeInfo2))
    return true;

  LongType minRank =
      shape::rank(shapeInfo1) < shape::rank(shapeInfo2) ? shape::rank(shapeInfo1) : shape::rank(shapeInfo2);

  for (LongType i = -1; i >= -minRank; --i)
    if (shape::sizeAt(shapeInfo1, i) != shape::sizeAt(shapeInfo2, i) && shape::sizeAt(shapeInfo1, i) != 1 &&
        shape::sizeAt(shapeInfo2, i) != 1)
      return false;

  return true;
}

bool ShapeUtils::areShapesBroadcastable(const std::vector<LongType>& shape1, const std::vector<LongType>& shape2) {
  const auto rank1 = shape1.size();
  const auto rank2 = shape2.size();

  // Scalars can be broadcast with anything
  if (rank1 == 0 || rank2 == 0)
    return true;

  const LongType minRank = rank1 < rank2 ? rank1 : rank2;

  for (LongType i = 1; i <= minRank; ++i)
    if (shape1[rank1 - i] != shape2[rank2 - i] && shape1[rank1 - i] != 1 && shape2[rank2 - i] != 1) return false;

  return true;
}

//////////////////////////////////////////////////////////////////////////
// check the possibility of broadcast operation, if true then return shapeInfo of resulting array
// if evalMinMax == false the array with larger rank has to be passed as first argument


bool ShapeUtils::evalBroadcastShapeInfo(NDArray& x, NDArray& y, const bool evalMinMax,
                                        LongType*& resultShapeInfo, memory::Workspace* workspace) {
  return evalBroadcastShapeInfo(x.shapeInfo(), y.shapeInfo(), evalMinMax, resultShapeInfo, workspace);
}



bool ShapeUtils::evalBroadcastShapeInfo( LongType* max,  LongType* min, const bool evalMinMax,
                                         LongType*& resultShapeInfo, memory::Workspace* workspace) {

  // Scalars can be broadcast with anything - result shape is the non-scalar
  if (shape::isScalar(max) || shape::isScalar(min)) {
    if (shape::isScalar(max) && shape::isScalar(min)) {
      // Both scalars - use max
      resultShapeInfo = ConstantShapeHelper::getInstance().createFromExisting(max);
    } else if (shape::isScalar(max)) {
      // max is scalar, min is not - result is min's shape
      resultShapeInfo = ConstantShapeHelper::getInstance().createFromExisting(min);
    } else {
      // min is scalar, max is not - result is max's shape
      resultShapeInfo = ConstantShapeHelper::getInstance().createFromExisting(max);
    }
    return true;
  }
  // Handle empty arrays early - if either input has a dimension of size 0, result should be empty
  bool maxEmpty = shape::isEmptyConst(max);
  bool minEmpty = shape::isEmptyConst(min);

  if (shape::shapeEquals(max, min)) {
    const int len = shape::shapeInfoLength(shape::rank(max));
    resultShapeInfo = new LongType[len];
    const auto constCast = const_cast<LongType*>(resultShapeInfo);

    for (int i = 0; i < len; i++) {
      constCast[i] = max[i];
    }
    resultShapeInfo = (ConstantShapeHelper::getInstance().bufferForShapeInfo(resultShapeInfo)->primary());
    return true;
  }

  // sometimes we have 1 and 2d vectors
  if (shape::isVector(min) && shape::isVector(max) && shape::length(min) == shape::length(max)) {
    if (shape::rank(min) > shape::rank(max)) {
      resultShapeInfo = ConstantShapeHelper::getInstance().createFromExisting(min);
      return true;
    }
    resultShapeInfo = ConstantShapeHelper::getInstance().createFromExisting(max);
    return true;
  }

  // check whether broadcast operation is possible for input arrays
  if (!areShapesBroadcastable(max, min)) {
    std::string errorMessage;
    errorMessage += "ShapeUtils::evalBroadcastShapeInfo: shapes are not broadcastable!\n";
    errorMessage += "Shape 1: ";
    errorMessage += ShapeUtils::shapeAsString(max);
    errorMessage += "\nShape 2: ";
    errorMessage += ShapeUtils::shapeAsString(min);
    errorMessage += "\n";
    THROW_EXCEPTION(errorMessage.c_str());
  }

  auto maxShapeInfo = max;
  auto minShapeInfo = min;
  if (evalMinMax && (shape::rank(max) < shape::rank(min))) {
    maxShapeInfo = min;
    minShapeInfo = max;
  }

  const auto maxRank = shape::rank(maxShapeInfo);
  const auto minRank = shape::rank(minShapeInfo);

  // evaluate shapeInfo for resulting array
  if (resultShapeInfo != nullptr)
    THROW_EXCEPTION(
        "std::runtime_error(ShapeUtils::evalBroadcastShapeInfo method: the input pointer on shapeInfo must be empty "
        "(=nullptr) !");

  LongType* tmpShapeInfo = nullptr;
  ALLOCATE(tmpShapeInfo, workspace, shape::shapeInfoLength(maxRank), sd::LongType);

  memcpy(tmpShapeInfo, maxShapeInfo, shape::shapeInfoByteLength(maxRank));

  // Handle dimension broadcasting - dimension size 0 should be preserved (empty arrays)
  // Compare dimensions from right to left (broadcasting semantics)
  for (LongType i = 0; i < minRank; ++i) {
    // Get dimensions from the end: -1 means last dim, -2 means second-to-last, etc.
    LongType maxDim = shape::sizeAt(maxShapeInfo, -1 - i);
    LongType minDim = shape::sizeAt(minShapeInfo, -1 - i);

    // If either dimension is 0, result should be 0 (empty array)
    if (maxDim == 0 || minDim == 0) {
      tmpShapeInfo[1 + maxRank - 1 - i] = 0;
    }
      // Otherwise follow standard broadcasting rules
    else if (maxDim < minDim) {
      tmpShapeInfo[1 + maxRank - 1 - i] = minDim;
    }
  }

  updateStridesAndType(tmpShapeInfo, DataTypeUtils::pickPairwiseResultType(maxShapeInfo, minShapeInfo),
                       shape::order(maxShapeInfo));

  if (maxEmpty || minEmpty) {
    ArrayOptions::setPropertyBit(tmpShapeInfo, ARRAY_EMPTY);
    memset(shape::stride(tmpShapeInfo), 0, shape::rank(tmpShapeInfo) * sizeof(LongType));
  }

  resultShapeInfo = (ConstantShapeHelper::getInstance().bufferForShapeInfo(tmpShapeInfo)->primary());
  delete[] tmpShapeInfo;
  return true;
}

//////////////////////////////////////////////////////////////////////////
// evaluate shapeInfo for resulting array from tile operation
LongType* ShapeUtils::evalTileShapeInfo(NDArray& arr, const std::vector<LongType>& reps,
                                        memory::Workspace* workspace) {
  // check whether reps contains at least one zero (then throw exception) or whether all elements in reps are unities
  // (then simply reshape or do nothing)
  LongType repsSize = reps.size();
  LongType product = 1;
  for (const auto& item : reps) product *= item;
  if (product == 0) THROW_EXCEPTION("NDArray::tile method: one of the elements in reps array is zero !");

  LongType rankOld = arr.rankOf();
  LongType diff = rankOld - repsSize;

  // evaluate new shapeInfo
  LongType* newShapeInfo = nullptr;
  if (diff < 0) {
    ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(repsSize), sd::LongType);
    newShapeInfo[0] = repsSize;  // set new rank
    for (LongType i = 1; i <= -diff; ++i)
      newShapeInfo[i] = 1;  // set unities to be new dimensions at left-hand side of newShapeInfo shape place
    memcpy(newShapeInfo + 1 - diff, arr.shapeInfo() + 1,
           rankOld * sizeof(LongType));  // copy old dimensions to the right-hand side of newShapeInfo shape place
    for (LongType i = 1; i <= repsSize; ++i)
      newShapeInfo[i] *= reps[i - 1];  // set new shape by multiplying old dimensions by corresponding numbers from reps
  } else {
    ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(rankOld), sd::LongType);
    memcpy(newShapeInfo, arr.shapeInfo(),
           shape::shapeInfoByteLength(rankOld));  // copy all elements of _shapeInfo to newShapeInfo
    for (LongType i = 1; i <= repsSize; ++i)
      newShapeInfo[rankOld + 1 - i] *=
          reps[repsSize - i];  // set new shape by multiplying old dimensions by corresponding numbers from reps
  }
  shape::updateStrides(newShapeInfo, arr.ordering(), false);
  ArrayOptions::setDataType(newShapeInfo, arr.dataType());

  auto ret = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo)->primary();
  RELEASE(newShapeInfo, workspace);
  return ret;
}

std::vector<LongType> ShapeUtils::pullShapeFromShapeInfo(const LongType* shapeInfo) {
  std::vector<LongType> shape(shape::rank(shapeInfo));
  LongType shapeSize = shape.size();

  for (LongType e = 0; e < shapeSize; e++) shape[e] = shape::shapeOf(shapeInfo)[e];

  return shape;
}

std::string ShapeUtils::shapeAsString(NDArray* array) {
  if (array->rankOf() == 0 && !array->isEmpty()) return "[0]";

  std::string result;

  result.append("[");
  for (LongType e = 0; e < array->rankOf(); e++) {
    result += flatbuffers::NumToString(array->sizeAt(e));
    if (e < array->rankOf() - 1) result.append(", ");
  }
  result.append("]");

  return result;
}

std::string ShapeUtils::shapeAsString(const std::vector<LongType>& shape) {
  std::string result;

  result.append("[");
  for (size_t e = 0; e < shape.size(); e++) {
    result += flatbuffers::NumToString(shape.at(e));
    if (e < shape.size() - 1) result.append(", ");
  }
  result.append("]");

  return result;
}

std::string ShapeUtils::shapeAsString(const LongType* shapeInfo) {
  if (shapeInfo == nullptr) THROW_EXCEPTION("ShapeUtils::shapeAsString method: input shapeInfo must not be nullptr !");

  if (shapeInfo[0] < 0 || shapeInfo[0] > SD_MAX_RANK) {
    THROW_EXCEPTION(
        "Shape info appears to be corrupt. Shape info[0] is less than 0 or greater than 32. Might have been "
        "deallocated.");
  }

  std::string result;

  result.append("[");
  for (LongType e = 0; e < shapeInfo[0]; e++) {
    result += flatbuffers::NumToString(shapeInfo[e + 1]);
    if (e < shapeInfo[0] - 1) result.append(", ");
  }
  result.append("]");

  return result;
}

std::string ShapeUtils::shapeInfoAsString(const LongType* shapeInfo) {
  if (!shapeInfo) THROW_EXCEPTION("ShapeUtils::shapeAsString method: input shapeInfo must not be nullptr !");

  std::string result;  // Stack allocation instead of heap

  LongType len = shape::shapeInfoLength(shapeInfo[0]);

  result.append("[");
  for (LongType e = 0; e < len; e++) {
    result.append(flatbuffers::NumToString(shapeInfo[e]));
    if (e < len - 1) result.append(", ");
  }
  result.append("]");

  return result;  // Return by value (move semantics will optimize)
}

std::string ShapeUtils::shapeAsString(const LongType rank, const LongType* shapeInfo) {
  if (!shapeInfo) THROW_EXCEPTION("ShapeUtils::shapeAsString method: input shapeInfo must not be nullptr !");

  std::string result;

  result.append("[");
  for (LongType e = 0; e < rank; e++) {
    result += flatbuffers::NumToString(shapeInfo[e]);
    if (e < rank - 1) result.append(", ");
  }
  result.append("]");

  return result;
}

//////////////////////////////////////////////////////////////////////////
std::vector<LongType> ShapeUtils::shapeAsVector(const LongType* shapeInfo) {
  if (!shapeInfo) THROW_EXCEPTION("ShapeUtils::shapeAsVector method: input shapeInfo must not be nullptr !");

  std::vector<LongType> vector(shapeInfo[0]);

  for (LongType e = 0; e < shapeInfo[0]; e++) vector[e] = shapeInfo[e + 1];

  return vector;
}

//////////////////////////////////////////////////////////////////////////
// evaluate shapeInfo for diagonal array which is made using input arr elements as diagonal
LongType* ShapeUtils::evalDiagShapeInfo(LongType* shapeInfoConst, memory::Workspace* workspace) {
  auto shapeInfo = const_cast<LongType*>(shapeInfoConst);

  const auto rank = shape::rank(shapeInfo);

  LongType* outputShapeInfo = nullptr;

  if (shape::isVector(shapeInfo) || shape::isScalar(shapeInfo)) {
    ALLOCATE(outputShapeInfo, workspace, shape::shapeInfoLength(2), sd::LongType);
    outputShapeInfo[0] = 2;
    outputShapeInfo[1] = outputShapeInfo[2] = shape::length(shapeInfo);
  } else {
    ALLOCATE(outputShapeInfo, workspace, shape::shapeInfoLength(2 * rank), sd::LongType);
    outputShapeInfo[0] = 2 * rank;
    for (LongType i = 1; i <= rank; ++i) outputShapeInfo[i] = outputShapeInfo[i + rank] = shapeInfo[i];
  }

  updateStridesAndType(outputShapeInfo, shapeInfo, shape::order(shapeInfo));
  auto nonConstShape = const_cast<LongType*>(outputShapeInfo);
  auto result = ConstantShapeHelper::getInstance().bufferForShapeInfo(nonConstShape);
  RELEASE(outputShapeInfo, workspace);
  return result->primary();
}

std::vector<LongType> ShapeUtils::evalBroadcastBackwardAxis(const LongType* operand, const LongType* result) {
  // rRank >= oRank always  !!
  const auto oRank = shape::rank(operand);
  const auto rRank = shape::rank(result);
  const auto diff = rRank - oRank;
  std::vector<LongType> axis;

  for (LongType i = 0; i < rRank; ++i)
    if (i < diff || shape::sizeAt(operand, i - diff) != shape::sizeAt(result, i)) axis.push_back(i);

  return axis;
}

////////////////////////////////////////////////////////////////////////////////
LongType* ShapeUtils::matrixProductShape(LongType* theFirstShape, LongType* theSecondShape,
                                         bool shouldTranspondFirst, bool shouldTranspondSecond, DataType dtype,
                                         memory::Workspace* workspace) {
  auto inA = theFirstShape;
  auto inB = theSecondShape;
  LongType* shape;
  ALLOCATE(shape, workspace, shape::shapeInfoLength(2), sd::LongType);

  LongType* tmpA = ShapeBuilders::copyShapeInfo(inA, true, workspace);
  LongType* tmpB = ShapeBuilders::copyShapeInfo(inB, true, workspace);

  if (shouldTranspondFirst) shape::transposeInplace(tmpA);

  if (shouldTranspondSecond) shape::transposeInplace(tmpB);

  if (shape::rank(tmpA) == 1 && shape::isMatrix(tmpB)) {
    // special case here
    shape[0] = 1;
    shape[1] = tmpB[2];
    LongType* newShape = ShapeBuilders::createShapeInfo(dtype, 'f', 2, shape, workspace, false);

    RELEASE(shape, workspace);
    RELEASE(tmpA, workspace);
    RELEASE(tmpB, workspace);

    return newShape;
  } else if (shape::isScalar(tmpA) && shape::isScalar(tmpB)) {
    // just scalar vs scalar
    shape[0] = 1;
    shape[1] = 1;
  } else if (shape::isMatrix(tmpA) && shape::isVector(tmpB)) {
    // gemv case
    if (shape::rank(tmpB) == 2) {
      shape[0] = tmpA[1];
      shape[1] = tmpB[2];
    } else {
      // we have new 1D shape here
      auto newShape = ShapeBuilders::createVectorShapeInfo(dtype, tmpA[1], workspace);

      RELEASE(shape, workspace);
      RELEASE(tmpA, workspace);
      RELEASE(tmpB, workspace);

      return newShape;
    }
  } else if ((shape::isMatrix(tmpA) && shape::isMatrix(tmpB)) || (shape::isVector(tmpA) && shape::isMatrix(tmpB)) ||
             (shape::isColumnVector(tmpA) && shape::isVector(tmpB))) {
    // gemm case
    shape[0] = tmpA[1];
    shape[1] = tmpB[2];
  } else if ((shape::isVector(tmpA) && shape::isScalar(tmpB)) || (shape::isScalar(tmpA) && shape::isVector(tmpB))) {
    // element-wise
    shape[0] = 1;
    shape[1] = (LongType)sd::math::sd_max<LongType>(shape::length(tmpA), shape::length(tmpB));
  } else if (shape::isRowVector(tmpA) && shape::isRowVector(tmpB)) {
    // dot case
    shape[0] = 1;
    shape[1] = 1;
  } else if (shape::isRowVector(tmpA) && shape::isColumnVector(tmpB)) {
    // dot case
    shape[0] = 1;
    shape[1] = 1;
  }

  auto newShape = ConstantShapeHelper::getInstance().createShapeInfo(dtype, 'f', 2, shape, -1);

  RELEASE(shape, workspace);

  RELEASE(tmpA, workspace);
  RELEASE(tmpB, workspace);
  return newShape;
}

////////////////////////////////////////////////////////////////////////////////
std::vector<LongType> ShapeUtils::composeShapeUsingDimsAndIdx(const std::vector<LongType>& dimsAndIdx) {
  auto size = dimsAndIdx.size();
  if (size % 2 != 0)
    THROW_EXCEPTION("ShapeUtils::composeShapeUsingDimsAndIdx static method: the size of input vector must be even !");

  size /= 2;

  std::vector<LongType> shape(size);
  LongType index;

  for (LongType i = 0; i < static_cast<LongType>(size); ++i) {
    index = dimsAndIdx[i + size];
    if (index > static_cast<LongType>(size - 1))
      THROW_EXCEPTION("ShapeUtils::composeShapeUsingDimsAndIdx static method: input index is too large !");
    shape[index] = dimsAndIdx[i];
  }

  return shape;
}

////////////////////////////////////////////////////////////////////////////////
std::vector<LongType> ShapeUtils::evalShapeForMatmul(const LongType* xShapeInfo, const LongType* yShapeInfo,
                                                     const bool transX, const bool transY) {
  const auto xRank = xShapeInfo[0];
  const auto yRank = yShapeInfo[0];

  const LongType x0Dim = transX ? xShapeInfo[xRank] : xShapeInfo[xRank - 1];
  const LongType y0Dim = transY ? yShapeInfo[yRank] : yShapeInfo[yRank - 1];
  const LongType x1Dim = transX ? xShapeInfo[xRank - 1] : xShapeInfo[xRank];
  const LongType y1Dim = transY ? yShapeInfo[yRank - 1] : yShapeInfo[yRank];

  if (xRank == 1 && yRank == 1) {  // dot case, output is scalar
    if (xShapeInfo[1] != yShapeInfo[1]) {
      sd_printf(
          "ShapeUtils::evalShapeForMatmul method: since input arrays are vectors they must have the same length, but "
          "got x length = %i, y length = %i !",
          xShapeInfo[1], yShapeInfo[1]);
      THROW_EXCEPTION("");
    }
    return std::vector<LongType>({});
  }

  if (xRank == 1 && yRank == 2) {  // vector x matrix, i.e. [4] x [4,5] = [5], output is vector
    if (xShapeInfo[1] != y0Dim) {
      sd_printf(
          "ShapeUtils::evalShapeForMatmul method: input arrays have inconsistent shapes for vector-matrix product: x "
          "%s, y %s !",
          ShapeUtils::shapeAsString(xShapeInfo).c_str(), ShapeUtils::shapeAsString(yShapeInfo).c_str());
      THROW_EXCEPTION("");
    }
    return std::vector<LongType>({y1Dim});
  }

  if (xRank == 2 && yRank == 1) {  // matrix x vector , i.e. [4,5] x [5] = [4], output is vector
    if (x1Dim != yShapeInfo[1]) {
      sd_printf(
          "ShapeUtils::evalShapeForMatmul method: input arrays have inconsistent shapes for vector-matrix product: x "
          "%s, y %s !",
          ShapeUtils::shapeAsString(xShapeInfo).c_str(), ShapeUtils::shapeAsString(yShapeInfo).c_str());
      THROW_EXCEPTION("");
    }
    return std::vector<LongType>({x0Dim});
  }

  // rest cases - usual 2Dx2D or batched mmul
  // Handle rank mismatch when one input has singleton leading dimensions
  // This supports ONNX Gemm patterns like [1,1,1,768] x [768,768] -> [1,1,1,768]
  if (xRank != yRank) {
    // Check if higher-rank input has all singleton leading dims that can be squeezed
    const LongType* higherRankInfo = xRank > yRank ? xShapeInfo : yShapeInfo;
    const LongType* lowerRankInfo = xRank > yRank ? yShapeInfo : xShapeInfo;
    const auto higherRank = xRank > yRank ? xRank : yRank;
    const auto lowerRank = xRank > yRank ? yRank : xRank;
    const auto rankDiff = higherRank - lowerRank;

    // Check if all leading dimensions are singletons (size 1)
    bool allLeadingSingleton = true;
    for (LongType i = 0; i < rankDiff; ++i) {
      if (higherRankInfo[i + 1] != 1) {
        allLeadingSingleton = false;
        break;
      }
    }

    if (allLeadingSingleton && lowerRank == 2) {
      // Can treat as 2D matmul with singleton batch dims preserved in output
      // For x having higher rank: x[1,1,...,M,K] @ y[K,N] -> [1,1,...,M,N]
      // For y having higher rank: x[M,K] @ y[1,1,...,K,N] -> [1,1,...,M,N]

      LongType outM, outN, xK, yK;

      if (xRank > yRank) {
        // x is higher rank [1,1,...,M,K], y is 2D [K,N]
        // Get M and K from x's last 2 dimensions
        const LongType xSecondLast = higherRankInfo[higherRank - 1];  // M (or K if transposed)
        const LongType xLast = higherRankInfo[higherRank];            // K (or M if transposed)
        // Get K and N from y
        const LongType yFirst = lowerRankInfo[1];   // K (or N if transposed)
        const LongType ySecond = lowerRankInfo[2];  // N (or K if transposed)

        if (transX) {
          outM = xLast;
          xK = xSecondLast;
        } else {
          outM = xSecondLast;
          xK = xLast;
        }

        if (transY) {
          yK = ySecond;
          outN = yFirst;
        } else {
          yK = yFirst;
          outN = ySecond;
        }
      } else {
        // y is higher rank [1,1,...,K,N], x is 2D [M,K]
        // Get M and K from x
        const LongType xFirst = lowerRankInfo[1];   // M (or K if transposed)
        const LongType xSecond = lowerRankInfo[2];  // K (or M if transposed)
        // Get K and N from y's last 2 dimensions
        const LongType ySecondLast = higherRankInfo[higherRank - 1];  // K (or N if transposed)
        const LongType yLast = higherRankInfo[higherRank];            // N (or K if transposed)

        if (transX) {
          outM = xSecond;
          xK = xFirst;
        } else {
          outM = xFirst;
          xK = xSecond;
        }

        if (transY) {
          yK = yLast;
          outN = ySecondLast;
        } else {
          yK = ySecondLast;
          outN = yLast;
        }
      }

      // Validate K dimensions match
      if (xK != yK) {
        std::string errorMessage;
        errorMessage += "ShapeUtils::evalShapeForMatmul static method: the dimensions of arrays are inconsistent: ";
        errorMessage += "xShape = " + shapeAsString(xShapeInfo) + ", ";
        errorMessage += "yShape = " + shapeAsString(yShapeInfo);
        errorMessage += " (xK=" + std::to_string(xK) + ", yK=" + std::to_string(yK) + ") ! \n";
        THROW_EXCEPTION(errorMessage.c_str());
      }

      std::vector<LongType> cShape;
      // Preserve leading singleton dimensions from the higher-rank input
      for (LongType i = 0; i < rankDiff; ++i) {
        cShape.push_back(1);
      }
      // Add the matrix dimensions [M, N]
      cShape.push_back(outM);
      cShape.push_back(outN);
      return cShape;
    } else {
      sd_printf(
          "ShapeUtils::evalShapeForMatmul static method: the ranks of arrays must be the same, but got xRank = %i and "
          "yRank = %i ! \n",
          xRank, yRank);
      THROW_EXCEPTION("");
    }
  }

  if (x1Dim != y0Dim) {
    std::string errorMessage;
    errorMessage += "ShapeUtils::evalShapeForMatmul static method: the dimensions of arrays are inconsistent: ";
    errorMessage += "xShape = " + shapeAsString(xShapeInfo) + ", ";
    errorMessage += "yShape = " + shapeAsString(yShapeInfo) + " ! \n";
    THROW_EXCEPTION(errorMessage.c_str());
  }

  for (LongType i = 0; i < xRank - 2; ++i)
    if (xShapeInfo[i + 1] != yShapeInfo[i + 1]) {
      std::string errorMessage;
      errorMessage += "ShapeUtils::evalShapeForMatmul static method: the dimensions of arrays are inconsistent: ";
      errorMessage += "xShape = " + shapeAsString(xShapeInfo) + ", ";
      errorMessage += "yShape = " + shapeAsString(yShapeInfo) + " ! \n";
      THROW_EXCEPTION(errorMessage.c_str());
    }

  std::vector<LongType> cShape;
  for(int i = 0; i < xRank  - 2; i++) {
    cShape.push_back(shape::sizeAt(xShapeInfo, i));
  }
  cShape.push_back(x0Dim);
  cShape.push_back(y1Dim);

  return cShape;
}

////////////////////////////////////////////////////////////////////////////////
LongType ShapeUtils::getNumOfSubArrs(const LongType* shapeInfo, const std::vector<LongType>& dimsToExclude) {
  LongType numOfSubArrs = 1;

  if (static_cast<sd::LongType>(dimsToExclude.size()) == shape::rank(shapeInfo) ||
      dimsToExclude.size() == 0)  // means there is only one sub-array and it coincides with whole array
    return numOfSubArrs;

  for (const auto& dim : dimsToExclude) numOfSubArrs *= shapeInfo[dim + 1];

  return numOfSubArrs;
}

////////////////////////////////////////////////////////////////////////////////
void ShapeUtils::updateStridesAndType(LongType* dest, const LongType* source, const char order) {
  shape::updateStrides(dest, order, false);
  dest[2 * dest[0] + 1] = 0;  // zero extra
  ArrayOptions::copyDataType(dest, source);
}

////////////////////////////////////////////////////////////////////////////////
void ShapeUtils::updateStridesAndType(LongType* dest, const DataType dtype, const char order) {
  shape::updateStrides(dest, order, true);
  ArrayOptions::setDataType(dest, dtype);
}

bool ShapeUtils::areShapesEqual(const LongType* shapeInfo, const std::vector<LongType>& shapeOnly) {
  LongType  rank = shape::rank(shapeInfo);
  if (rank != static_cast<sd::LongType>(shapeOnly.size())) {
    return false;
  }

  sd::LongType  *inputShapeOnly = shape::shapeOf(shapeInfo);
  for (LongType i = 0; i < rank; ++i) {
    if (inputShapeOnly[i] != shapeOnly[i]) {
      return false;
    }
  }

  return true;
}

////////////////////////////////////////////////////////////////////////////////
std::vector<LongType>* ShapeUtils::evalDimsForReduceOp(const LongType rank,
                                                       const std::vector<LongType>* dimsToExclude) {
  std::vector<LongType>* dims = evalDimsToExclude(rank, dimsToExclude->size(), dimsToExclude->data());
  std::vector<LongType>* output = new std::vector<LongType>(*dims);

  LongType dimsExcludeLen = static_cast<LongType>(dimsToExclude->size());
  for (LongType j = 0; j < dimsExcludeLen; j++) {
    LongType currElement = dimsToExclude->at(j);
    bool contains = false;
    for (size_t i = 0; i < output->size(); i++) {
      if (output->at(i) == currElement) {
        contains = true;
        break;
      } else {
        contains = false;
      }
    }

    bool elementLess = currElement < rank;
    if (!contains && elementLess) {
      output->push_back(dimsToExclude->at(j));
    }
  }

  delete dims;
  return output;
}

////////////////////////////////////////////////////////////////////////////////

}  // namespace sd