Hello,
I am trying to implement numpy.arrange
using MPSGraph primitives. Along the way I got to this code, which I would expect is correct, however the program crashes with information about type mismatch.
What am I doing wrong?
Code
let G = MPSGraph()
let length = G.constant(9, shape: [1], dataType: .int32)
let base = G.constant(1, shape: [1], dataType: .int32)
let template = G.broadcast(base, shapeTensor: length, name: nil)
let r = G.for(
numberOfIterations: length,
initialBodyArguments: [template],
// just return the argument
body: { i, args in [args[0]] },
name: nil
)
Error
-:22:11: error: 'scf.for' op types mismatch between 0th iter operand and defined value
-:22:11: note: see current operation: %19 = "scf.for"(%16, %17, %18, %12) ( {
^bb0(%arg1: index, %arg2: tensor<*xi32>): // no predecessors
%21 = "tensor.from_elements"(%arg1) : (index) -> tensor<1xindex>
%22 = "mps.cast"(%21) {resultElementType = i32} : (tensor<1xindex>) -> tensor<1xi32>
"scf.yield"(%arg2) : (tensor<*xi32>) -> ()
}) : (index, index, index, tensor<9xi32>) -> tensor<*xi32>
/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph_Sim/MetalPerformanceShadersGraph-2.0.22/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphExecutable.mm:1052: failed assertion `Error: MLIR pass manager failed'
P.S. Is there a way I can return differently-shaped tensor from each iteration? In other words:
body: { i, args in
[G.concatTensors([args[0], i], dimension: 0, name: nil)]
}