Type mismatch with MPSGraph.for

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)]
}

The problem is with broadcasted template. Tensor with "static" shape works OK.

It seems, that this is indeed impossible due to MLIR dialect for Structured Control Flow.

https://llvm.discourse.group/t/scf-whileop-type-inference-with-dynamic-shapes/2529

Type mismatch with MPSGraph.for
 
 
Q