#pragma clang diagnostic push
#pragma ide diagnostic ignored "cert-err58-cpp"

/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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), created on 24.11.2017
// @author Paul Dubs
//

#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_mean_pairwssqerr_loss)

#include <ops/declarable/CustomOperations.h>

#include <iostream>
#include <numeric>

namespace sd {
namespace ops {

//////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(mean_pairwssqerr_loss, 3, 1, false, 0, 1) {
  /*
   * Implementation of mean pairwise squared error loss
   *
   * For context on where this loss function may be useful see:
   *
   * Wei, Z., Zhang, J., Shen, X., Lin, Z., Mech, R., Hoai, M. and Samaras, D., 2018.
   * Good view hunting: learning photo composition from dense view pairs. In Proceedings of the IEEE Conference on
   * Computer Vision and Pattern Recognition (pp. 5437-5446).
   *
   * The paper defines the loss function as:
   *
   * L(y,q) = 1/((n*(n-1))/2) * (sum_(i,j=1..n,i!=j)((y_i - y_j) - (q_i - q_j))^2)
   *
   * with y: predictions, q: labels, n: length of y and q
   *
   * As creating those pairs is computationally expensive, we implement a mathematically equivalent function:
   *
   * L(y,q) = 4/(n*(n-1)) * (n * sum (y_i - q_i)^2 - (sum y_i - q_i)^2)
   *
   * This equivalency can be derived as:
   *
   * sum_(i,j=1..n,i!=j)((y_i - y_j) - (q_i - q_j))^2 = sum_(i,j=1..n,i!=j)((y_i - q_i) - (y_j - q_j))^2
   *
   * To simplify the following equations we use
   *
   * sum_(i,j=1..n,i!=j)(d_i - d_j)^2 = sum_(i,j=1..n,i!=j)(d_i^2 + d_j^2 - 2*d_i*d_j)
   *
   * Due to the pairings each element will appear as both d_i and d_j exactly n-1 times. This allows us to split the
   * sum:
   *
   * sum_(i,j=1..n,i!=j)(d_i^2 + d_j^2 - 2*d_i*d_j) = 2*(n-1)*sum d_i^2 - 2 * sum_(i,j=1..n,i!=j) d_i * d_j
   *                                                = 2*((n-1) * sum d_i^2 - sum_(i,j=1..n,i!=j) d_i * d_j)
   *
   * Now we use the following equivalency:
   *
   * (sum d_i)^2 = sum d_i^2 + sum_(i,j=1..n,i!=j) d_i * d_j
   *
   * This allows us to now use sum d_i^2 and (sum d_i)^2 as a quick way to calculate the sum:
   *
   * (n-1) * sum d_i^2 - sum_(i,j=1..n,i!=j) d_i * d_j = n * sum d_i^2 - (sum d_i)^2
   *
   * And by substituting it into the original definition we get:
   *
   * 1/((n*(n-1))/2) * 2*(n * sum d_i^2 - (sum d_i)^2)
   *
   * Which can be again simplified to
   *
   * 4/(n*(n-1)) * (n * sum d_i^2 - (sum d_i)^2)
   *
   * After substituting d_i back to (y_i - q_i) this results in the function that we actually implement.
   *
   */
  auto predictions = INPUT_VARIABLE(0);
  auto weights = INPUT_VARIABLE(1);
  auto labels = INPUT_VARIABLE(2);

  auto output = OUTPUT_VARIABLE(0);

  int reductionMode =
      INT_ARG(0);  // 0 - "none"; 1 - "weighted_sum";  2 - "weighted_mean";  3 - "weighted_sum_by_nonzero_weights"

  // input validation
  REQUIRE_TRUE(labels->isSameShape(predictions), 0,
               "MEAN_PAIRWSSQERR_LOSS OP: labels and predictions arrays must have the same shapes, but got %s and %s "
               "correspondingly !",
               ShapeUtils::shapeAsString(labels).c_str(), ShapeUtils::shapeAsString(predictions).c_str());
  // only 4 possible reduction modes exist
  REQUIRE_TRUE(reductionMode == 0 || reductionMode == 1 || reductionMode == 2 || reductionMode == 3, 0,
               "MEAN_PAIRWSSQERR_LOSS OP: reduction mode value is not acceptable, possible values are 0, 1, 2, 3, but "
               "got %i instead!",
               reductionMode);

  if (labels->rankOf() == 1) {  // If labels and predictions are of rank 1, it means that all data entries are 0-tensor
    // (scalar) so that the result of becomes always zero.
    *output = 0.;
    return Status::OK;
  }

  std::vector<LongType> zero;
  zero.push_back(0);
  std::vector<LongType> *reductionIdx = ShapeUtils::evalDimsToExclude(labels->rankOf(),1,zero.data());

  auto n = double(labels->sizeAt(1));
  
  // Compute diffs = predictions - labels
  NDArray* diffs_ptr = (*predictions) - (*labels);
  NDArray diffs = *diffs_ptr;
  delete diffs_ptr;

  // Compute sumOfSquares = sum(diffs^2)
  NDArray* diffsSquared = diffs * diffs;
  NDArray* sumOfSquares_ptr = diffsSquared->reduceAlongDimension(reduce::Sum, reductionIdx, true);
  delete diffsSquared;
  NDArray sumOfSquares = *sumOfSquares_ptr;
  delete sumOfSquares_ptr;

  // Compute squareOfSum = (sum(diffs))^2
  NDArray* squareOfSum_ptr = diffs.reduceAlongDimension(reduce::Sum, reductionIdx, true);
  NDArray squareOfSum = *squareOfSum_ptr;
  delete squareOfSum_ptr;
  squareOfSum.applyScalar(scalar::Pow, 2, &squareOfSum);

  delete reductionIdx;
  
  // Compute E = ((sumOfSquares * n) - squareOfSum) * (4 / (n * (n - 1)))
  NDArray* sumOfSquaresTimesN = sumOfSquares * n;
  NDArray* numerator = (*sumOfSquaresTimesN) - squareOfSum;
  delete sumOfSquaresTimesN;
  
  NDArray* E_ptr = (*numerator) * (4 / (n * (n - 1)));
  delete numerator;
  
  NDArray E = *E_ptr;
  delete E_ptr;

  // weights array can be single scalar or has the same rank as labels, and must be broadcastable to labels
  REQUIRE_TRUE(weights->isScalar() || weights->rankOf() == E.rankOf(), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: weights array should be scalar or have the same rank as results array, "
               "but got %i and %i correspondingly!",
               weights->rankOf(), E.rankOf());
  // check whether broadcast operation is possible for weights array
  REQUIRE_TRUE(weights->isScalar() || ShapeUtils::areShapesBroadcastable(*weights, E), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: shapes of weights and labels arrays should be broadcastable, but got "
               "weights = %s and results = %s instead!",
               ShapeUtils::shapeAsString(weights).c_str(), ShapeUtils::shapeAsString(&E).c_str());

  // perform weights broadcasting/tile to labels if needed
  auto weightsBroad = weights;
  if (!weights->isScalar() && !weights->isSameShape(E)) weightsBroad = new NDArray(weights->tileToShape(E.shapeInfo()));

  E *= *weightsBroad;

  switch (reductionMode) {
    case 0:  // 0 - "none", un-reduced weighted losses with the same shape as labels.
      output->assign(&E);
      break;

    case 1: {  // 1 - "weighted_sum", output is scalar and equal to sum of all elements of E array
      E.reduceNumber(reduce::Sum, output);
      break;
    }
    case 2: {  // 2 - "weighted_mean", output is scalar and equal to sum of all elements of E array divided by sum of
      // all elements of weightsBroad array
      NDArray* sumPtr = nullptr;
      if (weights->isScalar()) {
        NDArray* weightTimesLen = (*weights) * E.lengthOf();
        sumPtr = weightTimesLen;
      } else {
        sumPtr = weightsBroad->reduceNumber(reduce::Sum);
      }

      if (sumPtr->e<double>(0) == 0.) {
        (*output) = 0.;
      } else {
        NDArray* eSum = E.reduceNumber(reduce::Sum);
        NDArray* result = (*eSum) / (*sumPtr);
        delete eSum;
        output->assign(result);
        delete result;
      }
      delete sumPtr;
      break;
    }
    case 3: {  // 3 - "weighted_sum_by_nonzero_weights", output is scalar and equal to scalar sum of all elements of E
      // array divided by number of non-zero weights
      LongType numOfNonZeroWeights = 0;
      if (weights->isScalar()) {
        if (weights->e<double>(0) != 0.) numOfNonZeroWeights = E.lengthOf();
      } else {
        NDArray* countNonZero = weightsBroad->reduceNumber(reduce::CountNonZero);
        numOfNonZeroWeights = countNonZero->e<LongType>(0);
        delete countNonZero;
      }

      if (numOfNonZeroWeights == 0)
        (*output) = 0.;
      else {
        NDArray* eSum = E.reduceNumber(reduce::Sum);
        NDArray* result = (*eSum) / double(numOfNonZeroWeights);
        delete eSum;
        output->assign(result);
        delete result;
      }
      break;
    }
  }

  if (weightsBroad != weights) delete weightsBroad;

  return Status::OK;
}

//////////////////////////////////////////////////////////////////////////
DECLARE_TYPES(mean_pairwssqerr_loss) {
  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
}

//////////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(mean_pairwssqerr_loss) {
  auto predictionsShapeInfo = inputShape->at(0);
  auto weightsShapeInfo = inputShape->at(1);
  auto labelsShapeInfo = inputShape->at(2);

  REQUIRE_TRUE(shape::shapeEquals(labelsShapeInfo, predictionsShapeInfo), 0,
               "MEAN_PAIRWSSQERR_LOSS OP: labels and predictions arrays must have the same shapes, but got %s and %s "
               "correspondingly !",
               ShapeUtils::shapeAsString(labelsShapeInfo).c_str(),
               ShapeUtils::shapeAsString(predictionsShapeInfo).c_str());
  DataType outType = DataTypeUtils::pickFloatingType(ArrayOptions::dataType(predictionsShapeInfo));
  LongType  *outShapeInfo = nullptr;

  if (INT_ARG(0) != 0)  // in this case output is scalar
    outShapeInfo = ConstantShapeHelper::getInstance().scalarShapeInfo(outType);
  else {  // in this case output has the shape as labels and logits minus last dimension
    std::vector<LongType> dimensions = {-1};
    outShapeInfo = ShapeUtils::evalReduceShapeInfo(shape::order(predictionsShapeInfo), &dimensions, predictionsShapeInfo,
                                                   false, true, block.getWorkspace());

    // weights array can be single scalar or has the same rank as output, and must be broadcastable to output
    REQUIRE_TRUE(shape::isScalar(weightsShapeInfo) || shape::rank(weightsShapeInfo) == shape::rank(outShapeInfo), 0,
                 "MEAN_PAIRWSSQERR_LOSS OP: weights array should be scalar or have the same rank as output array, but "
                 "got %i and %i correspondingly!",
                 shape::rank(weightsShapeInfo), shape::rank(outShapeInfo));
    // check whether broadcast operation is possible for weights array
    REQUIRE_TRUE(
        shape::isScalar(weightsShapeInfo) || ShapeUtils::areShapesBroadcastable(weightsShapeInfo, outShapeInfo), 0,
        "MEAN_PAIRWSSQERR_LOSS OP: shapes of weights and output arrays should be broadcastable, but got weights = %s "
        "and output = %s instead!",
        ShapeUtils::shapeAsString(weightsShapeInfo).c_str(), ShapeUtils::shapeAsString(outShapeInfo).c_str());
  }

  return SHAPELIST(outShapeInfo);
}

//////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(mean_pairwssqerr_loss_grad, 3, 3, false, 0, 1) {
  auto predictions = INPUT_VARIABLE(0);
  auto weights = INPUT_VARIABLE(1);
  auto labels = INPUT_VARIABLE(2);

  auto dLdp = OUTPUT_VARIABLE(0);  // dL/dpredictions
  auto dLdw = OUTPUT_VARIABLE(1);  // dL/dweights
  auto dLdl = OUTPUT_VARIABLE(2);  // dL/dlabels

  int reductionMode =
      INT_ARG(0);  // 0 - "none"; 1 - "weighted_sum";  2 - "weighted_mean";  3 - "weighted_sum_by_nonzero_weights"
  // take into account Alex's proposition to treat "none" the same as "weighted_sum" mode when calculating gradients
  if (reductionMode == 0) reductionMode = 1;

  // inputs validation
  REQUIRE_TRUE(labels->isSameShape(predictions), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: labels and predictions arrays must have the same shapes, but got %s and "
               "%s correspondingly !",
               ShapeUtils::shapeAsString(labels).c_str(), ShapeUtils::shapeAsString(predictions).c_str());
  // only 4 possible reduction modes exist
  REQUIRE_TRUE(reductionMode == 0 || reductionMode == 1 || reductionMode == 2 || reductionMode == 3, 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: reduction mode value is not acceptable, possible values are 0, 1, 2, 3, "
               "but got %i instead!",
               reductionMode);

  auto n = double(labels->sizeAt(1));
  
  // Compute diffs = predictions - labels
  NDArray* diffs_ptr = (*predictions) - (*labels);
  NDArray diffs = *diffs_ptr;
  delete diffs_ptr;
  
  std::vector<LongType> dims2;
  dims2.push_back(0);
  std::vector<LongType> *reductionIdx = ShapeUtils::evalDimsToExclude(labels->rankOf(), 1,dims2.data());
  
  // Compute sumOfSquares
  NDArray* diffsSquared = diffs * diffs;
  NDArray* sumOfSquares_ptr = diffsSquared->reduceAlongDimension(reduce::Sum, reductionIdx, true);
  delete diffsSquared;
  NDArray sumOfSquares = *sumOfSquares_ptr;
  delete sumOfSquares_ptr;

  // Compute squareOfSum
  NDArray* squareOfSum_ptr = diffs.reduceAlongDimension(reduce::Sum, reductionIdx, true);
  NDArray squareOfSum = *squareOfSum_ptr;
  delete squareOfSum_ptr;
  squareOfSum.applyScalar(scalar::Pow, 2, &squareOfSum);

  // Compute E
  NDArray* sumOfSquaresTimesN = sumOfSquares * n;
  NDArray* eTerm1 = (*sumOfSquaresTimesN) - squareOfSum;
  delete sumOfSquaresTimesN;
  NDArray* E_ptr = (*eTerm1) * (4 / (n * (n - 1)));
  delete eTerm1;
  NDArray E = *E_ptr;
  delete E_ptr;

  // Compute gradients for predictions
  NDArray* sumPred_ptr = predictions->reduceAlongDimension(reduce::Sum, reductionIdx, true);
  NDArray sumPred = *sumPred_ptr;
  delete sumPred_ptr;
  
  NDArray* sumLabel_ptr = labels->reduceAlongDimension(reduce::Sum, reductionIdx, true);
  NDArray sumLabel = *sumLabel_ptr;
  delete sumLabel_ptr;

  delete reductionIdx;
  
  // Compute dLdp = ((diffs * n) - sumPred + sumLabel) * (8 / (n * (n - 1)))
  NDArray* diffsTimesN = diffs * n;
  NDArray* term1 = (*diffsTimesN) - sumPred;
  delete diffsTimesN;
  NDArray* term2 = (*term1) + sumLabel;
  delete term1;
  NDArray* dLdpTemp = (*term2) * (8 / (n * (n - 1)));
  delete term2;
  dLdp->assign(dLdpTemp);
  delete dLdpTemp;

  // weights array can be single scalar or has the same rank as labels, and must be broadcastable to labels
  REQUIRE_TRUE(weights->isScalar() || weights->rankOf() == E.rankOf(), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: weights array should be scalar or have the same rank as results array, "
               "but got %i and %i correspondingly!",
               weights->rankOf(), E.rankOf());
  // check whether broadcast operation is possible for weights array
  REQUIRE_TRUE(weights->isScalar() || ShapeUtils::areShapesBroadcastable(*weights, E), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: shapes of weights and labels arrays should be broadcastable, but got "
               "weights = %s and results = %s instead!",
               ShapeUtils::shapeAsString(weights).c_str(), ShapeUtils::shapeAsString(&E).c_str());

  // perform weights broadcasting/tile to labels if needed
  auto weightsBroad = weights;
  if (!weights->isScalar() && !weights->isSameShape(E)) weightsBroad = new NDArray(weights->tileToShape(E.shapeInfo()));

  switch (reductionMode) {
    case 1: {  // 1 - "none" and "weighted_sum", output is scalar and equal to sum of all elements of E array

      *dLdp *= *weightsBroad;

      if (weights->isScalar()) {
        NDArray* eSum = E.reduceNumber(reduce::Sum);
        dLdw->assign(eSum);
        delete eSum;
      } else if (weights != weightsBroad) {
        std::vector<LongType> axesToReduceAlong =
            ShapeUtils::evalBroadcastBackwardAxis(weights->shapeInfo(), weightsBroad->shapeInfo());
        E.reduceAlongDimension(reduce::Sum, dLdw, &axesToReduceAlong, true);
      } else
        dLdw->assign(&E);
      break;
    }
    case 2: {  // 2 - "weighted_mean", output is scalar and equal to sum of all elements of E array divided by sum of
      // all elements of weightsBroad array

      NDArray* sum = nullptr;
      if (weights->isScalar()) {
        NDArray* weightTimesLen = (*weights) * E.lengthOf();
        sum = weightTimesLen;
      } else {
        sum = weightsBroad->reduceNumber(reduce::Sum);
      }

      if (sum->e<double>(0) == 0.) {
        *dLdp = 0.;
        *dLdw = 0.;
      } else {
        NDArray* weightsDivSum = (*weightsBroad) / (*sum);
        *dLdp *= *weightsDivSum;
        delete weightsDivSum;

        if (weights->isScalar())
          *dLdw = 0.;
        else if (weights != weightsBroad) {
          std::vector<LongType> axesToReduceAlong =
              ShapeUtils::evalBroadcastBackwardAxis(weights->shapeInfo(), weightsBroad->shapeInfo());

          // Compute (E * sum - (E * weightsBroad).reduceNumber(Sum)) / (sum * sum)
          NDArray* ETimesSum = E * (*sum);
          NDArray* ETimesWeights = E * (*weightsBroad);
          NDArray* ETimesWeightsSum = ETimesWeights->reduceNumber(reduce::Sum);
          delete ETimesWeights;

          NDArray* numerator = (*ETimesSum) - (*ETimesWeightsSum);
          delete ETimesSum;
          delete ETimesWeightsSum;

          NDArray* sumSquared = (*sum) * (*sum);
          NDArray* result = (*numerator) / (*sumSquared);
          delete numerator;
          delete sumSquared;

          result->reduceAlongDimension(reduce::Sum, dLdw, &axesToReduceAlong, true);
          delete result;
        } else {
          // Compute (E * sum - (E * weightsBroad).reduceNumber(Sum)) / (sum * sum)
          NDArray* ETimesSum = E * (*sum);
          NDArray* ETimesWeights = E * (*weightsBroad);
          NDArray* ETimesWeightsSum = ETimesWeights->reduceNumber(reduce::Sum);
          delete ETimesWeights;

          NDArray* numerator = (*ETimesSum) - (*ETimesWeightsSum);
          delete ETimesSum;
          delete ETimesWeightsSum;

          NDArray* sumSquared = (*sum) * (*sum);
          NDArray* result = (*numerator) / (*sumSquared);
          delete numerator;
          delete sumSquared;

          dLdw->assign(result);
          delete result;
        }
      }
      delete sum;
      break;
    }
    case 3: {  // 3 - "weighted_sum_by_nonzero_weights", output is scalar and equal to scalar sum of all elements of E
      // array divided by number of non-zero weights

      LongType numOfNonZeroWeights = 0;
      if (weights->isScalar()) {
        if (weights->e<double>(0) != 0.) numOfNonZeroWeights = E.lengthOf();
      } else {
        NDArray* countNonZero = weightsBroad->reduceNumber(reduce::CountNonZero);
        numOfNonZeroWeights = countNonZero->e<LongType>(0);
        delete countNonZero;
      }

      if (numOfNonZeroWeights == 0) {
        *dLdp = 0.;
        *dLdw = 0.;
      } else {
        auto* numOfNonZeroWeightsScalar =
            NDArrayFactory::create(dLdw->dataType(), numOfNonZeroWeights, block.launchContext());

        if (weights->isScalar()) {
          NDArray* eSum = E.reduceNumber(reduce::Sum);
          NDArray* result = (*eSum) / double(numOfNonZeroWeights);
          delete eSum;
          dLdw->assign(result);
          delete result;
        } else if (weights != weightsBroad) {
          std::vector<LongType> axesToReduceAlong =
              ShapeUtils::evalBroadcastBackwardAxis(weights->shapeInfo(), weightsBroad->shapeInfo());
          E.reduceAlongDimension(reduce::Sum, dLdw, &axesToReduceAlong, true);
          *dLdw /= *numOfNonZeroWeightsScalar;
        } else {
          NDArray* eDivNum = E / (*numOfNonZeroWeightsScalar);
          dLdw->assign(eDivNum);
          delete eDivNum;
        }
        
        NDArray* weightsDivNum = (*weightsBroad) / (*numOfNonZeroWeightsScalar);
        NDArray temp = *weightsDivNum;
        delete weightsDivNum;
        *dLdp *= temp;
        
        delete numOfNonZeroWeightsScalar;
      }
      break;
    }
  }
  
  NDArray negDLdp = -(*dLdp);  // unary negation returns value
  dLdl->assign(&negDLdp);

  if (weightsBroad != weights) delete weightsBroad;

  return Status::OK;
}

DECLARE_TYPES(mean_pairwssqerr_loss_grad) {
  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
}

DECLARE_SHAPE_FN(mean_pairwssqerr_loss_grad) {
  auto predictionsShapeInfo = inputShape->at(0);
  auto weightsShapeInfo = inputShape->at(1);
  auto labelsShapeInfo = inputShape->at(2);

  // labels and predictions must have the same shapes
  REQUIRE_TRUE(shape::shapeEquals(labelsShapeInfo, predictionsShapeInfo), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: labels and predictions arrays must have the same shapes, but got %s and "
               "%s correspondingly !",
               ShapeUtils::shapeAsString(labelsShapeInfo).c_str(),
               ShapeUtils::shapeAsString(predictionsShapeInfo).c_str());
  // weights array can be single scalar or has the same rank as labels, and must be broadcastable to labels
  REQUIRE_TRUE(shape::isScalar(weightsShapeInfo) || shape::rank(weightsShapeInfo) == shape::rank(labelsShapeInfo), 0,
               "MEAN_PAIRWSSQERR_LOSS_GRAD OP: weights array should be scalar or have the same rank as labels array, "
               "but got %i and %i correspondingly!",
               shape::rank(weightsShapeInfo), shape::rank(labelsShapeInfo));
  // check whether broadcast operation is possible for weights array
  REQUIRE_TRUE(
      shape::isScalar(weightsShapeInfo) || ShapeUtils::areShapesBroadcastable(weightsShapeInfo, labelsShapeInfo), 0,
      "MEAN_PAIRWSSQERR_LOSS_GRAD OP: shapes of weights and labels arrays should be broadcastable, but got weights = "
      "%s and labels = %s instead!",
      ShapeUtils::shapeAsString(weightsShapeInfo).c_str(), ShapeUtils::shapeAsString(labelsShapeInfo).c_str());

  DataType outType = DataTypeUtils::pickFloatingType(ArrayOptions::dataType(predictionsShapeInfo));

  LongType *dLdpShapeInfo =
      ShapeBuilders::copyShapeInfoAndType(predictionsShapeInfo, outType, false, block.getWorkspace());
  LongType *dLdwShapeInfo = ShapeBuilders::copyShapeInfoAndType(weightsShapeInfo, outType, false, block.getWorkspace());
  LongType *dLdlShapeInfo =
      ShapeBuilders::copyShapeInfoAndType(labelsShapeInfo, outType, false, block.getWorkspace());

  return SHAPELIST(dLdpShapeInfo, dLdwShapeInfo, dLdlShapeInfo);
}
}  // namespace ops
}  // namespace sd

#endif
#pragma clang diagnostic pop