run() flattens outputs to 1-D by default. Use the shape= argument to reshape: out = kernel.run(din) # shape=(-1) -> 1-D out = kernel.run(din, shape=(4, 4)) # 2-D out1, out2 = kernel.run(in1, in2, shape=[(3,), (4,)]) # one per output The total element count must match the requested shape.