Skip to content
Snippets Groups Projects
Commit fd7061db authored by Bhatu's avatar Bhatu
Browse files

Fix reduce_mean/sum of tensors to 1 dimension.

Closes #96
parent 08421b2c
No related branches found
No related tags found
No related merge requests found
......@@ -1289,11 +1289,14 @@ class IRBuilderCSF(IRBuilderAST):
if node.keepdims == 1:
calculated_shape.append(1)
outputiters.append(IR.Int(0,32))
if calculated_shape == []:
calculated_shape = [1]
outputiters.append(IR.Int(0,32))
# perm will now be [ 1 ,2 ] + [ 0, 3]
perm.extend(reduced_dims)
loop_shape = [inputShape[perm[i]] for i in range(len(inputShape))]
outputShape = node.type.shape
assert(calculated_shape == outputShape)
assert calculated_shape == outputShape, "calculate shape:{} - real_shape: {}".format(calculated_shape, outputShape)
sumExpr = self.getTempVar()
sumExpr_decl = IR.Decl(sumExpr.idf, Type.Int())
......
......@@ -589,7 +589,11 @@ class TFNodesAST:
reductionAxesNodeName = inputsRef[1]
redAxesN = graph.__getitem__(reductionAxesNodeName)
redAxesT = redAxesN.getAttrVal("value").getTensor()
reductionAxesList = redAxesT.getContentAsValArr()
rank = redAxesT.getShapeRef().getRank()
if rank != 0:
reductionAxesList = redAxesT.getContentAsValArr()
else:
reductionAxesList = [redAxesT.getConstantVal()]
curNodeShapeLi = extraNodeInfoDict[curNode.getName()][0]
return (None, { curNode.getName() : AST.Reduce(AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]),
......@@ -609,7 +613,11 @@ class TFNodesAST:
reductionAxesNodeName = inputsRef[1]
redAxesN = graph.__getitem__(reductionAxesNodeName)
redAxesT = redAxesN.getAttrVal("value").getTensor()
reductionAxesList = redAxesT.getContentAsValArr()
rank = redAxesT.getShapeRef().getRank()
if rank != 0:
reductionAxesList = redAxesT.getContentAsValArr()
else:
reductionAxesList = [redAxesT.getConstantVal()]
curNodeShapeLi = extraNodeInfoDict[curNode.getName()][0]
return (None, { curNode.getName() : AST.Reduce(AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]),
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment