wehub/deeplearning4j--deeplearning4j
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libnd4j/include/ops/declarable/generic/blas/axpy.cpp
+64
View File
@@ -0,0 +1,64 @@
/* ******************************************************************************
*
*
* 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 raver119@gmail.com
//
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_axpy)
#include <ops/declarable/CustomOperations.h>
namespace sd {
namespace ops {
CONFIGURABLE_OP_IMPL(axpy, 2, 1, false, -2, 0) {
auto x = INPUT_VARIABLE(0);
auto y = INPUT_VARIABLE(1);
auto z = OUTPUT_VARIABLE(0);
REQUIRE_TRUE(x->isSameShape(y), 0, "Axpy: both arguments should have the same shape");
REQUIRE_TRUE(x->dataType() == y->dataType() && x->dataType() == z->dataType(), 0,
"Axpy: all arguments must have the same data type");
double a = 1.0;
if (block.width() > 2) {
auto alpha = INPUT_VARIABLE(2);
REQUIRE_TRUE(alpha->isScalar(), 0, "Axpy: alpha argument should be scalar or TArg");
} else if (block.getTArguments()->size() > 0) {
a = T_ARG(0);
}
ExtraArguments arguments({a});
y->applyPairwiseTransform(pairwise::Axpy, x, z, &arguments);
return Status::OK;
}
DECLARE_TYPES(axpy) {
getOpDescriptor()
->setAllowedInputTypes(0, {ALL_FLOATS})
->setAllowedInputTypes(1, {ALL_FLOATS})
->setAllowedOutputTypes(0, {ALL_FLOATS});
}
} // namespace ops
} // namespace sd
#endif
libnd4j/include/ops/declarable/generic/blas/batched_gemm.cpp
+355
View File
@@ -0,0 +1,355 @@
/* ******************************************************************************
*
*
* 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 raver119@gmail.com
// @author Adam Gibson
//
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_batched_gemm)
#include <ops/declarable/headers/blas.h>
#include <ops/declarable/helpers/batched_gemm.h>
namespace sd {
namespace ops {
CUSTOM_OP_IMPL(batched_gemm, -1, -1, false, 0, 2) {
// Only require 2 IArgs: transposeA, transposeB
// Everything else will be inferred from the input matrices
int transA = INT_ARG(0);
int transB = INT_ARG(1);
// Get alpha and beta
auto alpha = INPUT_VARIABLE(0);
auto beta = INPUT_VARIABLE(1);
// Calculate batch size from number of inputs
// Total inputs = alpha + beta + batchSize*A + batchSize*B
// So batchSize = (total - 2) / 2
int batchSize = (block.width() - 2) / 2;
REQUIRE_TRUE(batchSize > 0, 0, "BatchedGemm: Invalid batch size calculated: %d", batchSize);
REQUIRE_TRUE((block.width() - 2) % 2 == 0, 0, "BatchedGemm: Number of matrix inputs must be even");
// Get first matrices to infer dimensions
auto firstA = INPUT_VARIABLE(2);
auto firstB = INPUT_VARIABLE(2 + batchSize);
REQUIRE_TRUE(firstA->rankOf() == 2, 0, "BatchedGemm: A matrices must be rank 2");
REQUIRE_TRUE(firstB->rankOf() == 2, 0, "BatchedGemm: B matrices must be rank 2");
// Infer dimensions from first matrices
int M = transA ? firstA->sizeAt(1) : firstA->sizeAt(0);
int K = transA ? firstA->sizeAt(0) : firstA->sizeAt(1);
int N = transB ? firstB->sizeAt(0) : firstB->sizeAt(1);
int K_B = transB ? firstB->sizeAt(1) : firstB->sizeAt(0);
REQUIRE_TRUE(K == K_B, 0, "BatchedGemm: Incompatible dimensions - K from A is %d, K from B is %d", K, K_B);
// Infer leading dimensions
int ldA = firstA->sizeAt(0);
int ldB = firstB->sizeAt(0);
int ldC = M;
// Validate leading dimensions
if(transA == 0) {
int ldaComp = M > 1 ? M : 1;
if(ldA < ldaComp) THROW_EXCEPTION("LDA must be >= max(1,m) when transa == false");
} else {
int ldaComp = K > 1 ? K : 1;
if(ldA < ldaComp)
THROW_EXCEPTION("LDA must be >= max(1,k) when transa == true");
}
if(transB == 0) {
int ldBComp = K > 1 ? K : 1;
if(ldB < ldBComp) {
THROW_EXCEPTION("LDB must be >= max(1,k) when transb == false");
}
} else {
int ldbComp = N > 1 ? N : 1;
if(ldB < ldbComp)
THROW_EXCEPTION("LDB must be >= max(1,N) when transb == true");
}
int ldcComp = M > 1 ? M : 1;
if(ldC < ldcComp) {
THROW_EXCEPTION("LDC must be >= max(1,M)");
}
// Convert transpose flags to BLAS format
int transABlas = transA ? 112 : 111; // 112 = CblasTrans, 111 = CblasNoTrans
int transBBlas = transB ? 112 : 111;
REQUIRE_TRUE((transA == 0 || transA == 1) && (transB == 0 || transB == 1), 0,
"BatchedGemm: valid values for transA and transB are: 0/1 for NoTrans/Trans respectively")
REQUIRE_TRUE(M > 0 && N > 0 && K > 0 && ldA > 0 && ldB > 0 && ldC > 0, 0,
"BatchedGemm: Invalid dimensions M=%d, N=%d, K=%d, ldA=%d, ldB=%d, ldC=%d", M, N, K, ldA, ldB, ldC);
// Handle alpha and beta
NDArray *alphaInput = nullptr;
if(alpha->isScalar()) {
std::vector<sd::LongType> shape = {batchSize};
alphaInput = new NDArray('c',shape,alpha->dataType());
alphaInput->assign(alpha);
} else {
alphaInput = alpha;
}
NDArray *betaInput = nullptr;
if(beta->isScalar()) {
std::vector<LongType> shape = {batchSize};
betaInput = new NDArray('c',shape,beta->dataType());
betaInput->assign(beta);
} else {
betaInput = beta;
}
std::vector<NDArray*> vA(batchSize);
std::vector<NDArray*> vB(batchSize);
std::vector<NDArray*> vC(batchSize);
// Check data types - matrices should all match, alpha/beta can be different
auto firstMatrixType = firstA->dataType();
for (int e = 0; e < batchSize; e++) {
vA[e] = INPUT_VARIABLE(e + 2);
vB[e] = INPUT_VARIABLE(e + 2 + batchSize);
vC[e] = OUTPUT_VARIABLE(e);
REQUIRE_TRUE(firstMatrixType == vA[e]->dataType(), 0,
"BatchedGemm: all A matrices must have same data type");
REQUIRE_TRUE(firstMatrixType == vB[e]->dataType(), 0,
"BatchedGemm: all B matrices must have same data type");
REQUIRE_TRUE(firstMatrixType == vC[e]->dataType(), 0,
"BatchedGemm: all output matrices must have same data type as input matrices");
REQUIRE_TRUE(vA[e]->rankOf() == 2, 0, "BatchedGemm: batch %i, rank of A should be equal to 2", e);
REQUIRE_TRUE(vB[e]->rankOf() == 2, 0, "BatchedGemm: batch %i, rank of B should be equal to 2", e);
REQUIRE_TRUE(vC[e]->rankOf() == 2, 0, "BatchedGemm: batch %i, rank of C should be equal to 2", e);
// Verify dimensions are consistent across batch
int currM = transABlas == 111 ? vA[e]->sizeAt(0) : vA[e]->sizeAt(1);
int currK_A = transABlas == 111 ? vA[e]->sizeAt(1) : vA[e]->sizeAt(0);
int currK_B = transBBlas == 111 ? vB[e]->sizeAt(0) : vB[e]->sizeAt(1);
int currN = transBBlas == 111 ? vB[e]->sizeAt(1) : vB[e]->sizeAt(0);
REQUIRE_TRUE(currM == M, 0, "BatchedGemm: batch %i, inconsistent M dimension: expected %d, got %d", e, M, currM);
REQUIRE_TRUE(currK_A == K, 0, "BatchedGemm: batch %i, inconsistent K dimension in A: expected %d, got %d", e, K, currK_A);
REQUIRE_TRUE(currK_B == K, 0, "BatchedGemm: batch %i, inconsistent K dimension in B: expected %d, got %d", e, K, currK_B);
REQUIRE_TRUE(currN == N, 0, "BatchedGemm: batch %i, inconsistent N dimension: expected %d, got %d", e, N, currN);
}
helpers::bgemm(vA, vB, vC, alphaInput, betaInput, transABlas, transBBlas, M, N, K, ldA, ldB, ldC);
if(alphaInput != alpha) {
delete alphaInput;
}
if(betaInput != beta) {
delete betaInput;
}
return Status::OK;
};
DECLARE_SHAPE_FN(batched_gemm) {
// Only require 2 IArgs: transposeA, transposeB
int transA = INT_ARG(0);
int transB = INT_ARG(1);
// Calculate batch size from inputs
int batchSize = (block.width() - 2) / 2;
if (batchSize <= 0) {
auto shapeList = SHAPELIST();
shapeList->push_back(
ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(inputShape->at(0)), 'c', {1, 1}));
return shapeList;
}
// Get dimensions from first matrices
auto firstA = inputShape->at(2);
auto firstB = inputShape->at(2 + batchSize);
int M = transA ? shape::sizeAt(firstA, 1) : shape::sizeAt(firstA, 0);
int N = transB ? shape::sizeAt(firstB, 0) : shape::sizeAt(firstB, 1);
// Get data type from first matrix, not from alpha/beta
auto firstMatrixType = ArrayOptions::dataType(firstA);
// Check that all matrices have the same type (skip alpha and beta)
for (int e = 2; e < block.width(); e++) {
REQUIRE_TRUE(firstMatrixType == ArrayOptions::dataType(inputShape->at(e)), 0,
"BatchedGemm: all matrices must have same data type");
}
auto shapeList = SHAPELIST();
std::vector<LongType> shape({M, N});
for (int e = 0; e < batchSize; e++) {
auto newShape =
ConstantShapeHelper::getInstance().createShapeInfo(firstMatrixType, 'f', shape);
shapeList->push_back(newShape);
}
return shapeList;
}
DECLARE_TYPES(batched_gemm) {
getOpDescriptor()
->setAllowedInputTypes({ALL_FLOATS})
->setAllowedOutputTypes({ALL_FLOATS});
}
CUSTOM_OP_IMPL(batched_gemm_bp, -1, -1, false, 0, 2) {
// Only require 2 IArgs: transposeA, transposeB
int transA = INT_ARG(0);
int transB = INT_ARG(1);
// Calculate batch size
// Inputs: alpha, beta, batchSize*A, batchSize*B, batchSize*dLdC
// So batchSize = (total - 2) / 3
int batchSize = (block.width() - 2) / 3;
REQUIRE_TRUE(batchSize > 0, 0, "BatchedGemmBp: Invalid batch size calculated: %d", batchSize);
// Get dimensions from first matrices
auto firstA = INPUT_VARIABLE(2);
auto firstB = INPUT_VARIABLE(2 + batchSize);
auto firstDlDOut = INPUT_VARIABLE(2 + batchSize * 2);
int M = transA ? firstA->sizeAt(1) : firstA->sizeAt(0);
int K = transA ? firstA->sizeAt(0) : firstA->sizeAt(1);
int N = transB ? firstB->sizeAt(0) : firstB->sizeAt(1);
// Infer leading dimensions
int ldA = firstA->sizeAt(0);
int ldB = firstB->sizeAt(0);
int ldC = M;
std::vector<NDArray *> matricesA;
std::vector<NDArray *> matricesB;
std::vector<NDArray *> dlDOut;
std::vector<NDArray *> dldXOutputs;
std::vector<NDArray *> dldYOutputs;
for (int e = 0; e < batchSize; e++) {
matricesA.push_back(INPUT_VARIABLE(e + 2));
matricesB.push_back(INPUT_VARIABLE(e + 2 + batchSize));
dlDOut.push_back(INPUT_VARIABLE(e + 2 + batchSize * 2));
//alphas and betas are also set for outputs even though they're zero,every input needs a gradient
dldXOutputs.push_back(OUTPUT_VARIABLE(e + 2));
dldYOutputs.push_back(OUTPUT_VARIABLE(e + 2 + batchSize));
}
auto alpha = INPUT_VARIABLE(0);
NDArray *alphaInput = nullptr;
if(alpha->lengthOf() != batchSize) {
std::vector<sd::LongType> shape = {batchSize};
alphaInput = new NDArray('c',shape,alpha->dataType());
alphaInput->assign(alpha);
} else {
alphaInput = alpha;
}
auto beta = INPUT_VARIABLE(1);
NDArray *betaInput = nullptr;
if(beta->lengthOf() != batchSize) {
std::vector<sd::LongType> shape = {batchSize};
betaInput = new NDArray('c',shape,beta->dataType());
betaInput->assign(beta);
} else {
betaInput = beta;
}
// Convert transpose flags to BLAS format for helper function
int transABlas = transA ? 112 : 111;
int transBBlas = transB ? 112 : 111;
// First gradient computation: dL/dA = dL/dC @ B^T (or B if transB)
int transA1 = 0;
int transB1 = transB;
int M1 = dlDOut[0]->sizeAt(0);
int N1 = transB ? matricesB[0]->sizeAt(0) : matricesB[0]->sizeAt(1);
int k1 = dlDOut[0]->sizeAt(1);
int lda1 = dlDOut[0]->sizeAt(0);
int ldb1 = matricesB[0]->sizeAt(0);
int ldc1 = dldXOutputs[0]->sizeAt(0);
helpers::bgemm(dlDOut, matricesB, dldXOutputs, alphaInput, betaInput,
transA1 ? 112 : 111, transB1 ? 112 : 111, M1, N1, k1, lda1, ldb1, ldc1);
// Second gradient computation: dL/dB = A^T @ dL/dC (or A if transA)
int transA2 = transA ? 0 : 1;
int transB2 = 0;
int M2 = transA ? matricesA[0]->sizeAt(1) : matricesA[0]->sizeAt(0);
int N2 = dlDOut[0]->sizeAt(1);
int k2 = transA ? matricesA[0]->sizeAt(0) : matricesA[0]->sizeAt(1);
int lda2 = matricesA[0]->sizeAt(0);
int ldb2 = dlDOut[0]->sizeAt(0);
int ldc2 = dldYOutputs[0]->sizeAt(0);
helpers::bgemm(matricesA, dlDOut, dldYOutputs, alphaInput, betaInput,
transA2 ? 112 : 111, transB2 ? 112 : 111, M2, N2, k2, lda2, ldb2, ldc2);
if(alphaInput != alpha) {
delete alphaInput;
}
if(betaInput != beta) {
delete betaInput;
}
return Status::OK;
};
DECLARE_SHAPE_FN(batched_gemm_bp) {
// Calculate batch size
int batchSize = (block.width() - 2) / 3;
auto xConstant = CONSTANT(inputShape->at(2));
auto yConstant = CONSTANT(inputShape->at(2 + batchSize));
auto ret = SHAPELIST();
//alpha
ret->push_back(xConstant);
//beta
ret->push_back(yConstant);
for(int i = 0; i < batchSize; i++) {
ret->push_back(xConstant);
}
for(int i = 0; i < batchSize; i++) {
ret->push_back(yConstant);
}
return ret;
}
DECLARE_TYPES(batched_gemm_bp) {
getOpDescriptor()
->setAllowedInputTypes({ALL_FLOATS})
->setAllowedOutputTypes({ALL_FLOATS});
}
} // namespace ops
} // namespace sd
#endif
libnd4j/include/ops/declarable/generic/blas/matmul.cpp
+352
View File
@@ -0,0 +1,352 @@
/* ******************************************************************************
*
*
* 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 raver119@gmail.com, created on 07.10.2017.
// @author GS <sgazeos@gmail.com>, modified
// @author Yurii Shyrma (iuriish@yahoo.com), fully rewritten
//
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_matmul)
#include <helpers/MmulHelper.h>
#include <ops/declarable/CustomOperations.h>
namespace sd {
namespace ops {
//////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(matmul, 2, 1, false, 0, -2) {
auto x = INPUT_VARIABLE(0);
auto y = INPUT_VARIABLE(1);
auto z = OUTPUT_VARIABLE(0);
if(x->isEmpty() || y->isEmpty())
return Status::OK;
int iSize = (int)block.getIArguments()->size();
int transX = iSize > 0 ? INT_ARG(0) : 0;
int transY = iSize > 1 ? INT_ARG(1) : 0;
const int transZ = iSize > 2 ? INT_ARG(2) : 0;
// optional use alpha nad beta
iSize = (int)block.getTArguments()->size();
double alpha = iSize > 0 ? T_ARG(0) : 1.0;
double beta = iSize > 1 ? T_ARG(1) : 0.0;
if (transZ) {
x = INPUT_VARIABLE(1);
y = INPUT_VARIABLE(0);
bool temp = transX;
transX = !transY;
transY = !temp;
}
// Compute ranks AFTER potential transZ swap
const int xRank = x->rankOf();
const int yRank = y->rankOf();
const int zRank = z->rankOf();
const int xLastDim = transX ? -2 : -1;
const int yLastDim = transY ? -2 : -1;
const int xLastButOneDim = transX ? -1 : -2;
const int yLastButOneDim = transY ? -1 : -2;
// ******* input validation ******* //
REQUIRE_TRUE(xRank > 0 && yRank > 0, 0,
"MATMUL OP: input arrays must have rank bigger than 0 (should not be scalars), but got instead: x rank "
"= %i, y rank = %i !",
xRank, yRank);
// 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]
NDArray* xReshaped = nullptr;
NDArray* yReshaped = nullptr;
NDArray* zReshaped = nullptr;
if (xRank != yRank && xRank > 2 && yRank == 2) {
// Check if x has all singleton leading dims
bool allLeadingSingleton = true;
for (int i = 0; i < xRank - 2; ++i) {
if (x->sizeAt(i) != 1) {
allLeadingSingleton = false;
break;
}
}
if (allLeadingSingleton) {
// Reshape x from [1,1,...,M,K] to [M,K] for matmul
std::vector<LongType> newXShape = {x->sizeAt(-2), x->sizeAt(-1)};
xReshaped = new NDArray(x->reshape(x->ordering(), newXShape));
// Reshape z from [1,1,...,M,N] to [M,N]
std::vector<LongType> newZShape = {z->sizeAt(-2), z->sizeAt(-1)};
zReshaped = new NDArray(z->reshape(z->ordering(), newZShape));
x = xReshaped;
z = zReshaped;
}
} else if (xRank != yRank && yRank > 2 && xRank == 2) {
// Check if y has all singleton leading dims
bool allLeadingSingleton = true;
for (int i = 0; i < yRank - 2; ++i) {
if (y->sizeAt(i) != 1) {
allLeadingSingleton = false;
break;
}
}
if (allLeadingSingleton) {
// Reshape y from [1,1,...,K,N] to [K,N] for matmul
std::vector<LongType> newYShape = {y->sizeAt(-2), y->sizeAt(-1)};
yReshaped = new NDArray(y->reshape(y->ordering(), newYShape));
// Reshape z from [1,1,...,M,N] to [M,N]
std::vector<LongType> newZShape = {z->sizeAt(-2), z->sizeAt(-1)};
zReshaped = new NDArray(z->reshape(z->ordering(), newZShape));
y = yReshaped;
z = zReshaped;
}
}
// Update ranks after potential reshaping
const int xRankFinal = x->rankOf();
const int yRankFinal = y->rankOf();
const int zRankFinal = z->rankOf();
if (xRankFinal == 1 && yRankFinal == 1) { // dot case, output is scalar (or vector with length = 1)
REQUIRE_TRUE(x->lengthOf() == y->lengthOf(), 0,
"MATMUL OP: since input arrays are vectors they must have the same length, but got x length = %i, y "
"length = %i !",
x->lengthOf(), y->lengthOf());
} else if (xRankFinal == 1 && yRankFinal == 2) { // vector x matrix, i.e. [4] x [4,5] = [5], output is vector
REQUIRE_TRUE(x->lengthOf() == y->sizeAt(yLastButOneDim), 0,
"MATMUL OP: input arrays have inconsistent shapes for vector-matrix product: x %s, y %s !",
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str());
} else if (xRankFinal == 2 && yRankFinal == 1) { // matrix x vector , i.e. [4,5] x [5] = [4], output is vector
REQUIRE_TRUE(x->sizeAt(xLastDim) == y->lengthOf(), 0,
"MATMUL OP: input arrays have inconsistent shapes for matrix-vector product: x %s, y %s !",
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str());
} else {
REQUIRE_TRUE(xRankFinal == yRankFinal && yRankFinal == zRankFinal, 0,
"MATMUL OP: input and output arrays must have the same rank, but got instead: x rank = %i, y rank = "
"%i, z rank = %i !",
xRankFinal, yRankFinal, zRankFinal);
REQUIRE_TRUE(x->sizeAt(xLastDim) == y->sizeAt(yLastButOneDim) && x->sizeAt(xLastButOneDim) == z->sizeAt(-2) &&
y->sizeAt(yLastDim) == z->sizeAt(-1),
0, "MATMUL OP: input/output arrays have inconsistent shapes for matrix product: x %s, y %s, z %s !",
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str(),
ShapeUtils::shapeAsString(z).c_str());
if (xRankFinal > 2) // outer dims must be the same
for (int i = 0; i < xRankFinal - 2; ++i)
REQUIRE_TRUE(x->sizeAt(i) == y->sizeAt(i) && y->sizeAt(i) == z->sizeAt(i), 0,
"MATMUL OP: input/output arrays have inconsistent shapes for matrix product: x %s, y %s, z %s !",
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str(),
ShapeUtils::shapeAsString(z).c_str());
}
// ******* end of input validation ******* //
MmulHelper::matmul(x, y, z, transX, transY, alpha, beta, z);
// Clean up reshaped arrays
delete xReshaped;
delete yReshaped;
delete zReshaped;
return Status::OK;
}
DECLARE_SYN(mMul, matmul);
DECLARE_SYN(mmul, matmul);
DECLARE_SYN(gemm, matmul);
DECLARE_SYN(gemv, matmul);
DECLARE_SYN(dot, matmul);
//////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(matmul) {
auto xShapeInfo = inputShape->at(0);
auto yShapeInfo = inputShape->at(1);
const int iSize = (int)block.getIArguments()->size();
int transX = iSize > 0 ? INT_ARG(0) : 0;
int transY = iSize > 1 ? INT_ARG(1) : 0;
const int transZ = iSize > 2 ? INT_ARG(2) : 0;
if (transZ) {
xShapeInfo = inputShape->at(1);
yShapeInfo = inputShape->at(0);
bool temp = transX;
transX = !transY;
transY = !temp;
}
auto zShapeOnly = ShapeUtils::evalShapeForMatmul(xShapeInfo, yShapeInfo, transX, transY);
auto dtypeX = ArrayOptions::dataType(xShapeInfo);
auto dtypeY = ArrayOptions::dataType(yShapeInfo);
auto xOrder = shape::order(xShapeInfo);
auto yOrder = shape::order(yShapeInfo);
auto zOrder = xOrder == 'c' && yOrder == 'c' ? 'c' : 'f';
// we just pick the higher data type out of X and Y
auto dtypeZ = dtypeX > dtypeY ? dtypeX : dtypeY;
if(shape::isEmptyConst(xShapeInfo) || shape::isEmptyConst(yShapeInfo)) {
return SHAPELIST(ConstantShapeHelper::getInstance().emptyShapeInfoWithShape(ArrayOptions::dataType(xShapeInfo),zShapeOnly));
}
auto newShape = ConstantShapeHelper::getInstance().createShapeInfo(dtypeZ, zOrder, zShapeOnly);
return SHAPELIST(newShape);
}
//////////////////////////////////////////////////////////////////////
DECLARE_TYPES(matmul) {
getOpDescriptor()
->setAllowedInputTypes(0, {ALL_FLOATS, ALL_INTS})
->setAllowedInputTypes(1, {ALL_FLOATS, ALL_INTS})
->setAllowedOutputTypes(0, {ALL_FLOATS, ALL_INTS});
}
//////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(matmul_bp, 3, 2, false, 0, -2) {
auto x = INPUT_VARIABLE(0);
auto y = INPUT_VARIABLE(1);
auto eps = INPUT_VARIABLE(2);
auto dldx = OUTPUT_VARIABLE(0);
auto dldy = OUTPUT_VARIABLE(1);
int iSize = (int)block.getIArguments()->size();
int transX = iSize > 0 ? INT_ARG(0) : 0;
int transY = iSize > 1 ? INT_ARG(1) : 0;
const int transZ = iSize > 2 ? INT_ARG(2) : 0;
// optional use alpha nad beta
iSize = (int)block.getTArguments()->size();
double alpha = iSize > 0 ? T_ARG(0) : 1.0;
double beta = iSize > 1 ? T_ARG(1) : 0.0;
/*
In: x=[a,b], y=[b,c]
tX tY tZ x y z dz dLdx dLdy
F F F [a,b] [b,c] [a,c] [a,c] [a,c]*[b,c]T = [a,b] x*yT [a,b]T*[a,c] = [b,c] xT*y T F
F [b,a] [b,c] [a,c] [a,c] ([a,c]*[b,c]T)T = [b,a] (x*yT)T [b,a]*[a,c] = [b,c] x*y F T
F [a,b] [c,b] [a,c] [a,c] ([a,c]*[c,b]) = [a,b] x*y [a,b]T*[a,c] = [b,c] ->T xT*y T T
F [b,a] [c,b] [a,c] [a,c] ([a,c]*[c,b])T = [b,a] (x*y)T [b,a]*[a,c] = [b,c] ->T x*y F F
T [a,b] [b,c] [c,a] [c,a]
*/
// special case for scalar value
if (eps->isScalar()) {
if (x->isVector() && y->isVector()) {
if (x->isRowVector() && y->isRowVector()) {
float ySum = y->sumNumber().e<float>(0);
NDArray *dldxTemp = (*eps) * ySum;
dldx->assign(dldxTemp);
delete dldxTemp;
float xSum = x->sumNumber().e<float>(0);
NDArray *dldyTemp = (*eps) * xSum;
dldy->assign(dldyTemp);
delete dldyTemp;
} else if (x->isColumnVector() && y->isColumnVector()) {
float ySum = y->sumNumber().e<float>(0);
NDArray *dldxTemp = (*eps) * ySum;
dldx->assign(dldxTemp);
delete dldxTemp;
float xSum = x->sumNumber().e<float>(0);
NDArray *dldyTemp = (*eps) * xSum;
dldy->assign(dldyTemp);
delete dldyTemp;
} else {
NDArray *dldxTemp = (*eps) * (*y);
dldx->assign(dldxTemp);
delete dldxTemp;
NDArray *dldyTemp = (*eps) * (*x);
dldy->assign(dldyTemp);
delete dldyTemp;
}
} else {
// assign all ones to shape as baseline
auto alphaBetaBase = 1.0f;
if (alpha > 0.0f) {
alphaBetaBase *= alpha;
}
if (beta > 0.0f) {
alphaBetaBase += beta;
}
dldx->assign(alphaBetaBase);
dldy->assign(alphaBetaBase);
// match the dimensions for reduction for matrix multiply: columns on first input, rows on second input
// the dimensions should match the matching dimensions to compute proper gradients wrt each input
// core gradient for each is sum(input) * eps as scalar
std::vector<LongType> axesZero({0});
NDArray *xSum = x->reduceAlongDimension(reduce::Sum, &axesZero);
NDArray *xSumScaled = *xSum * (*eps);
std::vector<sd::LongType> xSumShape = {xSumScaled->lengthOf(), 1};
NDArray* xSumRow = xSumScaled->reshape(xSumScaled->ordering(), xSumShape);
std::vector<LongType> axes({1});
NDArray *ySum = y->reduceAlongDimension(reduce::Sum, &axes);
NDArray *ySumScaled = *ySum * (*eps);
std::vector<sd::LongType> ySumShape = {1, ySumScaled->lengthOf()};
NDArray* ySumRow = ySumScaled->reshape(ySumScaled->ordering(), ySumShape);
// execute proper multiplication: rows for first input, columns for second
dldx->mulRowVector(ySumRow, dldx);
dldy->muliColumnVector(xSumRow);
// FIXED: Proper cleanup - delete each allocated array once, add missing cleanup
delete xSumRow;
delete xSumScaled;
delete xSum;
delete ySumRow;
delete ySumScaled;
delete ySum;
}
return Status::OK;
}
matmul op;
op.execute({eps, y}, {dldx}, {alpha, beta}, {transZ, !transY, transX}, {});
op.execute({x, eps}, {dldy}, {alpha, beta}, {!transX, transZ, transY}, {});
return Status::OK;
}
//////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(matmul_bp) {
return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)));
}
//////////////////////////////////////////////////////////////////////
DECLARE_TYPES(matmul_bp) {
getOpDescriptor()
->setAllowedInputTypes(0, {ALL_FLOATS})
->setAllowedInputTypes(1, {ALL_FLOATS})
->setAllowedInputTypes(2, {ALL_FLOATS})
->setAllowedOutputTypes(0, {ALL_FLOATS})
->setAllowedOutputTypes(1, {ALL_FLOATS});
}
} // namespace ops
} // namespace sd
#endif
libnd4j/include/ops/declarable/generic/blas/tensormmul.cpp
+268
View File
@@ -0,0 +1,268 @@
/* ******************************************************************************
*
*
* 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 raver119@gmail.com
//
#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_tensormmul)
#include <helpers/MmulHelper.h>
#include <helpers/ShapeUtils.h>
#include <ops/declarable/CustomOperations.h>
#include <numeric>
namespace sd {
namespace ops {
////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(tensormmul, 2, 1, false, 0, -1) {
auto a = INPUT_VARIABLE(0);
auto b = INPUT_VARIABLE(1);
auto c = OUTPUT_VARIABLE(0);
REQUIRE_TRUE(a->dataType() == b->dataType(), 0, "tensormmul: A, B and C data types must be the same");
// building axes
LongType axe0_size = INT_ARG(0);
LongType axe1_size = INT_ARG(axe0_size + 1);
std::vector<LongType> axes_0(axe0_size), axes_1(axe1_size);
for (LongType e = 0; e < axe0_size; e++) axes_0[e] = INT_ARG(e + 1);
for (LongType e = 0; e < axe1_size; e++) axes_1[e] = INT_ARG(e + axe0_size + 2);
std::vector<sd::LongType> permuteC = {};
MmulHelper::tensorDot(a, b, c, axes_0, axes_1,permuteC);
return Status::OK;
}
DECLARE_SYN(tensordot, tensormmul);
////////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(tensormmul) {
auto aShapeInfo = inputShape->at(0);
auto bShapeInfo = inputShape->at(1);
REQUIRE_TRUE(ArrayOptions::dataType(aShapeInfo) == ArrayOptions::dataType(bShapeInfo), 0,
"tensormmul: A and B data types must be the same");
// building axes
LongType axe0_size = INT_ARG(0);
LongType axe1_size = INT_ARG(axe0_size + 1);
std::vector<LongType> axes_0(axe0_size), axes_1(axe1_size);
for (LongType e = 0; e < axe0_size; e++) axes_0[e] = INT_ARG(e + 1);
for (LongType e = 0; e < axe1_size; e++) axes_1[e] = INT_ARG(e + axe0_size + 2);
sd_verbose("axe0: %i; axe1: %i;\n", axes_0.size(), axes_1.size());
// evaluate shapes
std::vector<LongType> permutAt, permutBt;
std::vector<LongType> shapeAt, shapeBt;
auto outShape =
ShapeUtils::evalShapeForTensorDot(aShapeInfo, bShapeInfo, axes_0, axes_1, permutAt, permutBt,
shapeAt, shapeBt);
auto desc = new ShapeDescriptor(ArrayOptions::dataType(aShapeInfo), 'c', outShape);
auto result = SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(desc));
delete desc;
return result;
}
////////////////////////////////////////////////////////////////////////
DECLARE_TYPES(tensormmul) {
getOpDescriptor()
->setAllowedInputTypes(0, {FLOAT32, DOUBLE, HALF})
->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
->setAllowedInputTypes(2, {FLOAT32, DOUBLE, HALF})
->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF});
}
// Comparator for sorting indices vector based on comparison of array values
struct IndexComparator
{
const std::vector<LongType>& array;
IndexComparator(const std::vector<LongType>& arr): array(arr) {}
bool operator() (LongType i1, LongType i2)
{
return array[i1] < array[i2];
}
};
std::vector<LongType> argsort(const std::vector<LongType>& array)
{
std::vector<LongType> indices(array.size());
for (size_t i = 0; i < array.size(); ++i) indices[i] = i;
std::sort(indices.begin(), indices.end(), IndexComparator(array));
return indices;
}
////////////////////////////////////////////////////////////////////////
CUSTOM_OP_IMPL(tensormmul_bp, 4, 2, false, 0, -1) {
auto A = INPUT_VARIABLE(0);
auto B = INPUT_VARIABLE(1);
auto C = INPUT_VARIABLE(2);
auto dC = INPUT_VARIABLE(3);
auto originalDC = dC;
//scalar case, tile value to be whatever the c value is. common when directly attached to the loss
if(dC->isScalar()) {
auto newVec = const_cast<NDArray *>(C);
auto* newShapeVec = newVec->getShapeAsVector();
dC = new NDArray('c', *newShapeVec, dC->dataType(), dC->getContext());
delete newShapeVec;
}
auto gradA = OUTPUT_VARIABLE(0);
auto gradB = OUTPUT_VARIABLE(1);
LongType axe0_size = INT_ARG(0);
LongType axe1_size = INT_ARG(axe0_size + 1);
std::vector<LongType> axes0Sum(axe0_size), axes1Sum(axe1_size);
//find the passed in axes for the feed forward
for (LongType e = 0; e < axe0_size; e++) axes0Sum[e] = INT_ARG(e + 1);
for (LongType e = 0; e < axe1_size; e++) axes1Sum[e] = INT_ARG(e + axe0_size + 2);
auto Arank = A->rankOf();
auto Brank = B->rankOf();
auto dCrank = dC->rankOf();
//part of the permtue axes before matrix multiply happens
std::vector<LongType> axes_a_grad;
for (LongType i = 0; i < Arank; ++i)
axes_a_grad.push_back(i);
for (size_t i = 0; i < axes0Sum.size(); ++i)
axes_a_grad.erase(std::remove(axes_a_grad.begin(), axes_a_grad.end(), axes0Sum[i]), axes_a_grad.end());
//part of matrix multiply axes before matrix multiply happens
std::vector<LongType> axes_b_grad;
for (LongType i = 0; i < Brank; ++i)
axes_b_grad.push_back(i);
for (size_t i = 0; i < axes1Sum.size(); ++i)
axes_b_grad.erase(std::remove(axes_b_grad.begin(), axes_b_grad.end(), axes1Sum[i]), axes_b_grad.end());
//used for post result permute to reshape result to be expected output
std::vector<LongType> grad_a_axes;
grad_a_axes.insert(grad_a_axes.end(), axes_a_grad.begin(), axes_a_grad.end());
grad_a_axes.insert(grad_a_axes.end(), axes1Sum.begin(), axes1Sum.end());
//used for post result permute to reshape result to be expected output
std::vector<LongType> grad_b_axes;
grad_b_axes.insert(grad_b_axes.end(), axes0Sum.begin(), axes0Sum.end());
grad_b_axes.insert(grad_b_axes.end(), axes_b_grad.begin(), axes_b_grad.end());
LongType starting = dCrank - axes_a_grad.size();
std::vector<LongType> axes_a_gradA;
for (LongType i = starting; i < dCrank; i++) {
axes_a_gradA.push_back(i);
}
std::vector<LongType> axes_b_gradA;
for (size_t i = 0; i < axes_b_grad.size(); i++) {
axes_b_gradA.push_back(i);
}
std::vector<LongType> axes_a_gradB;
for (size_t i = 0; i < axes_a_grad.size(); i++) {
axes_a_gradB.push_back(i);
}
LongType start = dCrank - axes_a_gradA.size();
std::vector<LongType> axes_b_gradB;
for (LongType i = start; i < dCrank; i++) {
axes_b_gradB.push_back(i);
}
//create final axes before for matrix multiply
std::vector<LongType> aPermuteAxesBefore;
aPermuteAxesBefore.insert(aPermuteAxesBefore.end(), axes_a_grad.begin(), axes_a_grad.end());
aPermuteAxesBefore.insert(aPermuteAxesBefore.end(), axes0Sum.begin(), axes0Sum.end());
//create final axes before for matrix multiply
std::vector<LongType> bPermuteAxesBefore;
bPermuteAxesBefore.insert(bPermuteAxesBefore.end(), axes_b_grad.begin(), axes_b_grad.end());
bPermuteAxesBefore.insert(bPermuteAxesBefore.end(), axes1Sum.begin(), axes1Sum.end());
auto aPermArgsAfter = argsort(grad_a_axes);
auto bPermArgsAfter = argsort(grad_b_axes);
auto newA = A->permute(aPermuteAxesBefore, false, false);
std::vector<LongType> empty;
auto newB = B->permute(bPermuteAxesBefore, false, false);
//perform the actual matrix multiplication
MmulHelper::tensorDot2(dC, newB, gradA, axes_a_gradA, axes_b_gradA, empty, empty, aPermArgsAfter, gradA);
MmulHelper::tensorDot2(newA, dC, gradB, axes_a_gradB, axes_b_gradB, empty, empty, bPermArgsAfter, gradB);
// FIXED: permute() with copyToNewBuff=false returns view - only delete if not view
if (newA != nullptr && !newA->isView()) {
delete newA;
}
if (newB != nullptr && !newB->isView()) {
delete newB;
}
if(dC != originalDC) {
delete dC;
}
return Status::OK;
}
////////////////////////////////////////////////////////////////////////
DECLARE_SHAPE_FN(tensormmul_bp) {
auto aShapeInfo = inputShape->at(0);
auto bShapeInfo = inputShape->at(1);
auto cShapeInfo = inputShape->at(2);
auto dLShapeInfo = inputShape->at(3);
REQUIRE_TRUE((ArrayOptions::dataType(aShapeInfo) == ArrayOptions::dataType(bShapeInfo) &&
(ArrayOptions::dataType(dLShapeInfo) == ArrayOptions::dataType(aShapeInfo))),
0, "tensormmul_bp: A, B and dLdC data types must be the same");
return SHAPELIST(CONSTANT(aShapeInfo), CONSTANT(bShapeInfo));
}
////////////////////////////////////////////////////////////////////////
DECLARE_TYPES(tensormmul_bp) {
getOpDescriptor()
->setAllowedInputTypes(0, {FLOAT32, DOUBLE, HALF}) // maybe better ALL_FLOATS
->setAllowedInputTypes(1, {FLOAT32, DOUBLE, HALF})
->setAllowedInputTypes(2, {FLOAT32, DOUBLE, HALF})
->setAllowedOutputTypes(0, {FLOAT32, DOUBLE, HALF})
->setAllowedOutputTypes(1, {FLOAT32, DOUBLE, HALF});
}
} // namespace ops
} // namespace sd
#endif