/* ******************************************************************************
 *
 *
 * 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 <cublas_v2.h>
#include <cuda_runtime.h>
#include <cusolverDn.h>
#include <exceptions/cuda_exception.h>
#include <helpers/PointersManager.h>
#include <helpers/ShapeUtils.h>
#include <helpers/svd.h>
#include <system/op_boilerplate.h>


namespace sd {
namespace ops {
namespace helpers {

// FIXME -> we should optimize these helpers for the case when input matrices have c order (perform transpositions
// appropriately)

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

//////////////////////////////////////////////////////////////////////////
static void svdQR(LaunchContext* context, NDArray* A, NDArray* S, NDArray* U, NDArray* VT, const bool fullUV,
                  const bool calcUV) {
  // since cusa api cusolverDnDgesvd/cusolverDnSgesvd have following constrain on input matrix A: A_rows >= A_columns &&
  // A_order = 'f' we make this function to have deal with 2 valid cases only: 1) A_rows >= A_columns and A_corder = 'f'
  // 2) A_rows <= A_columns and A_corder = 'c'    - int this case perform transposition to get f order
  // if 1) or 2) are not met then throw exception

  // A  [m, n]
  // S  [n]
  // U  [m, m] or [m, n] if fullUV = false and m > n
  // VT [n, n] or [m, n] if fullUV = false and m < n

  if (A->rankOf() != 2) THROW_EXCEPTION("svdQR: rank of A array is not equal 2 !");

  auto m = A->sizeAt(0);
  auto n = A->sizeAt(1);
  const int minDim = m < n ? m : n;
  const char orderA = A->ordering();

  if (m < n) THROW_EXCEPTION("svdQR: due to cuda api input constrains given shape of A array are not valid !");

  if (std::vector<LongType>({minDim}) != S->getShapeAsVector())
    THROW_EXCEPTION("svdQR: wrong shape of S array !");

  if (calcUV) {
    if (fullUV && std::vector<LongType>({m, m}) != U->getShapeAsVector()) {
      THROW_EXCEPTION("svdQR: wrong shape of U array !");
    } else if (!fullUV && std::vector<LongType>({m, minDim}) != U->getShapeAsVector()) {
      THROW_EXCEPTION("svdQR: wrong shape of U array !");
    }

    if (fullUV && std::vector<LongType>({n, n}) != VT->getShapeAsVector()) {
      THROW_EXCEPTION("svdQR: wrong shape of VT array !");
    }
    else if (!fullUV && std::vector<LongType>({minDim, n}) != VT->getShapeAsVector()) {
      THROW_EXCEPTION("svdQR: wrong shape of VT array !");
    }
  }

  NDArray* pA = const_cast<NDArray*>(A);
  NDArray* pS = S;
  NDArray* pU = U;
  NDArray* pVT = VT;

  std::vector<NDArray*> toDelete;

  if (pA->ordering() == 'c') {
    pA = A->dup('f');
    toDelete.push_back(pA);
  }

    pS =S->dup('f');
    toDelete.push_back(pS);


  if (calcUV) {
    if (pU->ordering() == 'c') {
      pU =U->dup('f');
      toDelete.push_back(pU);
    }

    if (pVT->ordering() == 'c') {
      pVT = VT->dup('f');
      toDelete.push_back(pVT);
    }
  }

  std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());

  // create cusolverDn handle
  cusolverDnHandle_t* handle = (cusolverDnHandle_t*)context->getCusolverHandle();  // nullptr;
  if (handle == nullptr) throw cuda_exception::build("svdQR: cuda failed !", -1);

  // stream
  auto status = cusolverDnSetStream(*handle, *context->getCudaStream());
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdQR: cuda failed !", status);

  // query working space of SVD
  int lwork = 0;
  if (A->dataType() == DOUBLE)
    status = cusolverDnDgesvd_bufferSize(*handle, m, n, &lwork);
  else if (A->dataType() == FLOAT32)
    status = cusolverDnSgesvd_bufferSize(*handle, m, n, &lwork);
  else
    THROW_EXCEPTION("svdQR: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdQR: cuda failed !", status);

  // allocate memory for dWork
  void* dWork = nullptr;
  cudaError_t status2 = cudaMalloc((void**)&dWork, A->sizeOfT() * lwork);
  if (status2 != cudaSuccess) throw cuda_exception::build("svdQR: cuda failed !", status2);

  signed char jobu, jobvt;

  if (calcUV) {
    if (fullUV)
      jobu = jobvt = 'A';
    else
      jobu = jobvt = 'S';
  } else {
    jobu = jobvt = 'N';
  }

  int* devInfo = nullptr;
  void* rWork = nullptr;

  int lda(m), ldu, ldvt;

  if (calcUV) {
    ldu = pU->sizeAt(0);
    ldvt = pVT->sizeAt(0);
  }

  PointersManager manager(context, "svdQR");

  NDArray::prepareSpecialUse({pS, pU, pVT}, {pA});

  // choose appropriate cuda gemm api depending on data types
  if (A->dataType() == DOUBLE) {
    status = cusolverDnDgesvd(*handle, jobu, jobvt, m, n, reinterpret_cast<double*>(pA->specialBuffer()), lda,
                              reinterpret_cast<double*>(pS->specialBuffer()),
                              calcUV ? reinterpret_cast<double*>(pU->specialBuffer()) : nullptr, ldu,
                              calcUV ? reinterpret_cast<double*>(pVT->specialBuffer()) : nullptr, ldvt,
                              reinterpret_cast<double*>(dWork), lwork, reinterpret_cast<double*>(rWork), devInfo);
  } else if (A->dataType() == FLOAT32) {
    status = cusolverDnSgesvd(*handle, jobu, jobvt, m, n, reinterpret_cast<float*>(pA->specialBuffer()), lda,
                              reinterpret_cast<float*>(pS->specialBuffer()),
                              calcUV ? reinterpret_cast<float*>(pU->specialBuffer()) : nullptr, ldu,
                              calcUV ? reinterpret_cast<float*>(pVT->specialBuffer()) : nullptr, ldvt,
                              reinterpret_cast<float*>(dWork), lwork, reinterpret_cast<float*>(rWork), devInfo);
  } else
    THROW_EXCEPTION("svdQR: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdQR: cuda failed !", status);

  manager.synchronize();

  NDArray::registerSpecialUse({pS, pU, pVT}, {pA});

  S->assign(pS);

  if (calcUV) {
    U->assign(pU);
    VT->assign(pVT);
  }

  //for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];

  // if (devInfo) cudaFree(devInfo);
  // if (dWork) cudaFree(dWork);
  // if (rWork) cudaFree(rWork);

}

//////////////////////////////////////////////////////////////////////////
static void svdJcb(LaunchContext* context, NDArray* A, NDArray* S, NDArray* U, NDArray* V, const bool fullUV,
                   const bool calcUV) {
  // A [m, n]
  // S [n]
  // U [m, m] or [m, n] if fullUV = false and m > n
  // V [n, n] or [n, m] if fullUV = false and m < n

  if (A->rankOf() != 2) THROW_EXCEPTION("svdJcb: rank of A array is not equal 2 !");

  int m = A->sizeAt(0);
  int n = A->sizeAt(1);
  const int minDim = m < n ? m : n;

  if (std::vector<LongType>({minDim}) != S->getShapeAsVector()) THROW_EXCEPTION("svdJcb: wrong shape of S array !");

  if (fullUV && U != nullptr && std::vector<LongType>({m, m}) != U->getShapeAsVector()) {
    THROW_EXCEPTION("svdJcb: wrong shape of U array !");
  } else if (!fullUV && U != nullptr && std::vector<LongType>({m, minDim}) != U->getShapeAsVector()) {
    THROW_EXCEPTION("svdJcb: wrong shape of U array !");
  }
  if (fullUV && V != nullptr && std::vector<LongType>({n, n}) != V->getShapeAsVector()) {
    THROW_EXCEPTION("svdJcb: wrong shape of V array !");
  } else if (!fullUV && V != nullptr && std::vector<LongType>({n, minDim}) != V->getShapeAsVector()) {
    THROW_EXCEPTION("svdJcb: wrong shape of V array !");
  }

  NDArray* pA = const_cast<NDArray*>(A);

  const bool aForder = m == 1 || A->strideAt(0) == 1;
  const bool aCorder = n == 1 || A->strideAt(1) == 1;

  const bool transA = !aForder && aCorder;
  const bool dupA = !aForder && !aCorder;

  std::vector<NDArray*> toDelete;

  if (dupA) {
    pA = A->dup('f');
    toDelete.push_back(pA);
  }

  NDArray* pS = S;

    pS = S->dup('f');
    toDelete.push_back(pS);


  NDArray *pU(nullptr), *pV(nullptr);

  int lda = transA ? pA->strideAt(0) : pA->strideAt(1);
  int ldu(transA ? n : m), ldv(transA ? m : n);
  bool uForder(true), vForder(true);


  if (calcUV) {
    pU = transA ? V : U;
    pV = transA ? U : V;

    uForder = pU->sizeAt(0) == 1 || pU->strideAt(0) == 1;
    vForder = pV->sizeAt(0) == 1 || pV->strideAt(0) == 1;

    if (!uForder) {
      pU = pU->dup('f');
      toDelete.push_back(pU);
    }

    if (!vForder) {
      pV = pV->dup('f');
      toDelete.push_back(pV);
    }

    ldu = pU->strideAt(1);
    ldv = pV->strideAt(1);
  }

  std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());

  // create cusolverDn handle
  cusolverDnHandle_t* handle = (cusolverDnHandle_t*)context->getCusolverHandle();
  if (handle == nullptr) throw cuda_exception::build("svdJcb: cuda failed !", -1);

  // stream
  auto status = cusolverDnSetStream(*handle, *context->getCudaStream());
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);

  // set parameters
  gesvdjInfo_t gesvdjParams = nullptr;
  status = cusolverDnCreateGesvdjInfo(&gesvdjParams);
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);
  status = cusolverDnXgesvdjSetTolerance(gesvdjParams, 1.e-7);  // tolerance
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);
  status = cusolverDnXgesvdjSetMaxSweeps(gesvdjParams, 15);  // max_sweeps
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);

  int* devInfo = nullptr;
  const cusolverEigMode_t jobz = calcUV ? CUSOLVER_EIG_MODE_VECTOR : CUSOLVER_EIG_MODE_NOVECTOR;
  const int econ = !fullUV;

  if (transA) math::sd_swap<int>(m, n);

  // *** avoid bug in cuda API ***
  void* nullPtr = nullptr;
  NDArray* arrToAvoidBugInAPI = nullptr;
  if (!calcUV && m != n) {
    int maxDim = m > n ? m : n;
    std::vector<LongType> shape = {maxDim, maxDim};
    arrToAvoidBugInAPI = new NDArray('c', shape, pA->dataType(), context);
    nullPtr = arrToAvoidBugInAPI->specialBuffer();
  }
  // ******************

  NDArray::prepareSpecialUse({pS, pU, pV}, {pA});

  // query working space of SVD
  int lwork = 0;
  if (A->dataType() == DOUBLE)
    status = cusolverDnDgesvdj_bufferSize(
        *handle, jobz, econ, m, n, reinterpret_cast<double*>(pA->specialBuffer()), lda,
        reinterpret_cast<double*>(pS->specialBuffer()),
        calcUV ? reinterpret_cast<double*>(pU->specialBuffer()) : reinterpret_cast<double*>(nullPtr), ldu,
        calcUV ? reinterpret_cast<double*>(pV->specialBuffer()) : reinterpret_cast<double*>(nullPtr), ldv, &lwork,
        gesvdjParams);
  else if (A->dataType() == FLOAT32)
    status = cusolverDnSgesvdj_bufferSize(
        *handle, jobz, econ, m, n, reinterpret_cast<float*>(pA->specialBuffer()), lda,
        reinterpret_cast<float*>(pS->specialBuffer()),
        calcUV ? reinterpret_cast<float*>(pU->specialBuffer()) : reinterpret_cast<float*>(nullPtr), ldu,
        calcUV ? reinterpret_cast<float*>(pV->specialBuffer()) : reinterpret_cast<float*>(nullPtr), ldv, &lwork,
        gesvdjParams);
  else
    THROW_EXCEPTION("svdJcb: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);

  // allocate memory dWork
  void* dWork = nullptr;
  auto status2 = cudaMalloc((void**)&dWork, A->sizeOfT() * lwork);
  if (status2 != cudaSuccess) throw cuda_exception::build("svdJcb: cuda failed !", status2);

  PointersManager manager(context, "svdJcb");

  // choose appropriate cuda gemm api depending on data types
  if (A->dataType() == DOUBLE) {
    status = cusolverDnDgesvdj(
        *handle, jobz, econ, m, n, reinterpret_cast<double*>(pA->specialBuffer()), lda,
        reinterpret_cast<double*>(pS->specialBuffer()),
        calcUV ? reinterpret_cast<double*>(pU->specialBuffer()) : reinterpret_cast<double*>(nullPtr), ldu,
        calcUV ? reinterpret_cast<double*>(pV->specialBuffer()) : reinterpret_cast<double*>(nullPtr), ldv,
        reinterpret_cast<double*>(dWork), lwork, devInfo, gesvdjParams);
  } else if (A->dataType() == FLOAT32) {
    status = cusolverDnSgesvdj(
        *handle, jobz, econ, m, n, reinterpret_cast<float*>(pA->specialBuffer()), lda,
        reinterpret_cast<float*>(pS->specialBuffer()),
        calcUV ? reinterpret_cast<float*>(pU->specialBuffer()) : reinterpret_cast<float*>(nullPtr), ldu,
        calcUV ? reinterpret_cast<float*>(pV->specialBuffer()) : reinterpret_cast<float*>(nullPtr), ldv,
        reinterpret_cast<float*>(dWork), lwork, devInfo, gesvdjParams);
  } else
    THROW_EXCEPTION("svdJcb: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdJcb: cuda failed !", status);

  manager.synchronize();

  NDArray::registerSpecialUse({pS, pU, pV}, {pA});

   S->assign(pS);

  if (calcUV) {
    if (!uForder) U->assign(transA ? pV : pU);
    if (!vForder) V->assign(transA ? pU : pV);
  }

  if (!calcUV && m != n) delete arrToAvoidBugInAPI;

  for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];

  if (devInfo) cudaFree(devInfo);
  if (dWork) cudaFree(dWork);
  if (gesvdjParams) cusolverDnDestroyGesvdjInfo(gesvdjParams);


}

//////////////////////////////////////////////////////////////////////////
static void svdBatched(LaunchContext* context, NDArray* A, NDArray* S, NDArray* U, NDArray* V,
                       const bool fullUV, const bool calcUV) {
  // A [..., m, n]
  // S [..., n]
  // U [..., m, m] or [..., m, n] if fullUV = false and m > n
  // V [..., n, n] or [..., n, m] if fullUV = false and m < n

  auto m = A->sizeAt(-2);
  auto n = A->sizeAt(-1);
  const int minDim = m < n ? m : n;
  const LongType bS = A->lengthOf() / (m * n);

  if (m > 32 || n > 32) THROW_EXCEPTION("svdBatched: numbers of rows and columns should be <= 32 !");

  if (minDim != S->sizeAt(-1)) THROW_EXCEPTION("svdBatched: wrong shape of S array !");

  if (calcUV) {
    if (U->sizeAt(-2) != m) THROW_EXCEPTION("svdBatched: wrong shape of U array !");
    if (U->sizeAt(-1) != (fullUV ? m : minDim)) THROW_EXCEPTION("svdBatched: wrong shape of U array !");
    if (U->lengthOf() / (U->sizeAt(-2) * U->sizeAt(-1)) != bS)
      THROW_EXCEPTION("svdBatched: wrong shape of U array !");

    if (V->sizeAt(-2) != n) THROW_EXCEPTION("svdBatched: wrong shape of V array !");
    if (V->sizeAt(-1) != (fullUV ? n : minDim)) THROW_EXCEPTION("svdBatched: wrong shape of V array !");
    if (V->lengthOf() / (V->sizeAt(-2) * V->sizeAt(-1)) != bS)
      THROW_EXCEPTION("svdBatched: wrong shape of V array !");
  }

  NDArray* pA = const_cast<NDArray*>(A);
  NDArray* pS = S;
  NDArray* pU = U;
  NDArray* pV = V;

  std::vector<NDArray*> toDelete;

  if (pA->ordering() == 'c') {
    pA = A->dup('f');
    toDelete.push_back(pA);
  }

    pS = S->dup('f');
    toDelete.push_back(pS);


  if (calcUV) {
    if (pU->ordering() == 'c') {
      pU = U->dup('f');
      toDelete.push_back(pU);
    }

    if (pV->ordering() == 'c') {
      pV = V->dup('f');  // dup() already returns NDArray*
      toDelete.push_back(pV);
    }
  }

  // create cusolverDn handle
  cusolverDnHandle_t handle = nullptr;
  cusolverStatus_t status = cusolverDnCreate(&handle);
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);

  // stream
  status = cusolverDnSetStream(handle, *context->getCudaStream());
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);

  // set parameters
  gesvdjInfo_t gesvdjParams = nullptr;
  status = cusolverDnCreateGesvdjInfo(&gesvdjParams);
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);
  status = cusolverDnXgesvdjSetTolerance(gesvdjParams, 1.e-7);  // tolerance
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);
  status = cusolverDnXgesvdjSetMaxSweeps(gesvdjParams, 15);  // max_sweeps
  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);

  // devInfo
  int* devInfo = nullptr;
  auto status2 = cudaMalloc((void**)&devInfo, sizeof(LongType) * bS);
  if (status2 != cudaSuccess) throw cuda_exception::build("svdBatched: cuda failed !", status2);
  status2 = cudaDeviceSynchronize();
  if (status2 != cudaSuccess) throw cuda_exception::build("svdJcb: cuda failed !", status2);

  const cusolverEigMode_t jobz = calcUV ? CUSOLVER_EIG_MODE_VECTOR : CUSOLVER_EIG_MODE_NOVECTOR;

  int lda(m), ldu, ldv;

  if (calcUV) {
    ldu = pU->sizeAt(-2);
    ldv = pV->sizeAt(-2);
  }

  // Ak (i,j) = A[i + 5*j + 25*k]

  // query working space of SVD
  int lwork = 0;
  if (A->dataType() == DOUBLE)
    status = cusolverDnDgesvdjBatched_bufferSize(handle, jobz, m, n, reinterpret_cast<double*>(pA->specialBuffer()),
                                                 lda, reinterpret_cast<double*>(pS->specialBuffer()),
                                                 calcUV ? reinterpret_cast<double*>(pU->specialBuffer()) : nullptr, ldu,
                                                 calcUV ? reinterpret_cast<double*>(pV->specialBuffer()) : nullptr, ldv,
                                                 &lwork, gesvdjParams, bS);
  else if (A->dataType() == FLOAT32)
    status = cusolverDnSgesvdjBatched_bufferSize(
        handle, jobz, m, n, reinterpret_cast<float*>(pA->specialBuffer()), lda,
        reinterpret_cast<float*>(pS->specialBuffer()), calcUV ? reinterpret_cast<float*>(pU->specialBuffer()) : nullptr,
        ldu, calcUV ? reinterpret_cast<float*>(pV->specialBuffer()) : nullptr, ldv, &lwork, gesvdjParams, bS);
  else
    THROW_EXCEPTION("svdBatched: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);

  // allocate memory dWork
  void* dWork = nullptr;
  status2 = cudaMalloc((void**)&dWork, A->sizeOfT() * lwork);
  if (status2 != cudaSuccess) throw cuda_exception::build("svdBatched: cuda failed !", status2);
  status2 = cudaDeviceSynchronize();
  if (status2 != cudaSuccess) throw cuda_exception::build("svdBatched: cuda failed !", status2);

  PointersManager manager(context, "svdBatched");

  NDArray::prepareSpecialUse({pS, pU, pV}, {pA});

  // choose appropriate cuda gemm api depending on data types
  if (A->dataType() == DOUBLE) {
    status = cusolverDnDgesvdjBatched(handle, jobz, m, n, reinterpret_cast<double*>(pA->specialBuffer()), lda,
                                      reinterpret_cast<double*>(pS->specialBuffer()),
                                      calcUV ? reinterpret_cast<double*>(pU->specialBuffer()) : nullptr, ldu,
                                      calcUV ? reinterpret_cast<double*>(pV->specialBuffer()) : nullptr, ldv,
                                      reinterpret_cast<double*>(dWork), lwork, devInfo, gesvdjParams, bS);
  } else if (A->dataType() == FLOAT32) {
    status = cusolverDnSgesvdjBatched(handle, jobz, m, n, reinterpret_cast<float*>(pA->specialBuffer()), lda,
                                      reinterpret_cast<float*>(pS->specialBuffer()),
                                      calcUV ? reinterpret_cast<float*>(pU->specialBuffer()) : nullptr, ldu,
                                      calcUV ? reinterpret_cast<float*>(pV->specialBuffer()) : nullptr, ldv,
                                      reinterpret_cast<float*>(dWork), lwork, devInfo, gesvdjParams, bS);
  } else
    THROW_EXCEPTION("svdBatched: given data type is unsupported !");

  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("svdBatched: cuda failed !", status);

  manager.synchronize();

  NDArray::registerSpecialUse({pS, pU, pV}, {pA});

  S->assign(pS);

  if (calcUV) {
    U->assign(pU);
    V->assign(pV);
  }

  for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];

  if (devInfo) cudaFree(devInfo);
  if (dWork) cudaFree(dWork);
  if (handle) cusolverDnDestroy(handle);
  if (gesvdjParams) cusolverDnDestroyGesvdjInfo(gesvdjParams);

}

////////////////////////////////////////////////////////////////////
void svd(LaunchContext* context, NDArray* x, const std::vector<NDArray*>& outArrs, const bool fullUV,
         const bool calcUV, const int switchNum) {
  NDArray* S = outArrs[0];
  NDArray* U = outArrs[1];
  NDArray* V = outArrs[2];

  NDArray::prepareSpecialUse({S, U, V}, {x});

  if (x->rankOf() == 2) {
    svdJcb(context, x, S, U, V, fullUV, calcUV);
  } else {
    ResultSet *tadsU(nullptr), *tadsV(nullptr);

    auto tadsX = x->allTensorsAlongDimension({x->rankOf() - 2, x->rankOf() - 1});
    auto tadsS = S->allTensorsAlongDimension({S->rankOf() - 1});

    if (calcUV) {
      tadsU = new ResultSet(U->allTensorsAlongDimension({U->rankOf() - 2, U->rankOf() - 1}));
      tadsV = new ResultSet(V->allTensorsAlongDimension({V->rankOf() - 2, V->rankOf() - 1}));
    }

    for (int i = 0; i < tadsX.size(); ++i)
      svdJcb(context, tadsX.at(i), tadsS.at(i), calcUV ? tadsU->at(i) : nullptr, calcUV ? tadsV->at(i) : nullptr,
             fullUV, calcUV);

    if (calcUV) {
      delete tadsU;
      delete tadsV;
    }
  }

  NDArray::registerSpecialUse({S, U, V}, {x});
}

}  // namespace helpers
}  // namespace ops
}  // namespace sd