Gather: simplify expression for constant scalar indices
@@ -87,6 +87,10 @@ public class Gather extends IntermediateOperation {
addSliceDimension(dataSliceDimensions, dataType.dimensions().get(i).name(), i);
}
+ if (indicesType.rank() == 0 && indices.isConstant()) {
+ ExpressionNode indexExpression = new ConstantNode(new DoubleValue(indices.getConstantValue().get().asDouble()));
+ addSliceDimension(dataSliceDimensions, dataType.dimensions().get(axis).name(), indexExpression);
+ } else {
List<Slice.DimensionValue<Reference>> indicesSliceDimensions = new ArrayList<>();
for (int i = 0; i < indicesType.rank(); ++i) {
addSliceDimension(indicesSliceDimensions, indicesType.dimensions().get(i).name(), axis + i);
@@ -94,12 +98,13 @@ public class Gather extends IntermediateOperation {
ExpressionNode sliceExpression = createSliceExpression(indicesSliceDimensions, indicesFunctionName);
ExpressionNode indexExpression = createIndexExpression(dataType, sliceExpression);
addSliceDimension(dataSliceDimensions, dataType.dimensions().get(axis).name(), indexExpression);
+ }
for (int i = axis + 1; i < dataType.rank(); ++i) {
addSliceDimension(dataSliceDimensions, dataType.dimensions().get(i).name(), i + indicesType.rank() - 1);
}
- sliceExpression = createSliceExpression(dataSliceDimensions, dataFunctionName);
+ ExpressionNode sliceExpression = createSliceExpression(dataSliceDimensions, dataFunctionName);
return Generate.bound(type.type(), wrapScalar(sliceExpression));
}
