Dynamic multidimensional output for TensorFlow Lite in Kotlin

Question

I have a TF lite model, which has multiple input and output. I ran it with:

   tflite.runForMultipleInputsOutputs(arrayOf(buffer), output)

In this case, output is:

    val output = HashMap<Int, Any>()
    output[1] = Array(1) { FloatArray(100) }
    output[3] = IntArray(1)
    output[4] = Array(1) { Array(100) { FloatArray(4) } }
    output[5] = Array(1) { FloatArray(100) }

Output is hardcoded but I want to create it dynamically, using tflite.getOutputTensor(i).shape() & tflite.getOutputTensor(i).dataType()

I created the next transformation function:

1. Call:

val outputBuffer = HashMap<Int, Any>()
   for (i in 0 until tflite.outputTensorCount) {
      val shape = tflite.getOutputTensor(i).shape()
      val dataType = tflite.getOutputTensor(i).dataType()
      outputBuffer[i] = getOutputTensor(shape, dataType)
}

2. getOutputTensor content

   private fun getOutputTensor(size: IntArray, type: DataType): Any {
        val last = size.last()
        var result = when (type) {
            DataType.FLOAT32 -> FloatArray(last)
            DataType.INT32, DataType.UINT8, DataType.INT8 -> IntArray(last)
            DataType.INT64 -> LongArray(last)
            DataType.BOOL -> BooleanArray(last)
            DataType.STRING -> arrayOfNulls<String>(last)
        }
        if (size.size == 1) {
            return result
        }
        for (i in size.size - 2 downTo 0) {
            result = Array(size[i]) { result }
        }
        return result
    }

I start by creating the last typed array (Int, Float, String. etc) and then go from tail to start and put it in the Array.

Exception: My function return an array of the wrong type, while the last component is correct, all parent arrays convert to Object

Hardcoded	My Function
	.

DataType error: cannot resolve DataType of [Ljava.lang.Object;

Question:

How to create a dynamic multidimensional array with primitive type in Koltin?

UPDATE: Possible solution:

One of the options is to replace array allocation with ByteBuffer, like:

val outputBuffer = HashMap<Int, Buffer>()
...
val shape = tflite.getOutputTensor(i).numElements()
val dataType = tflite.getOutputTensor(i).dataType()
getOutputTensor(shape, dataType, outputBuffer, index)
...
private fun getOutputTensor(
        size: Int,
        type: DataType,
        output: HashMap<Int, Buffer>,
        index: Int
    ) {
        when (type) {
            DataType.FLOAT32 -> {
                val buffer: FloatBuffer = FloatBuffer.allocate(size)
                buffer.order()
                output[index] = buffer
            }
            DataType.INT32 -> {
                val buffer = IntBuffer.allocate(size)
                buffer.order()
                output[index] = buffer
            }
            DataType.UINT8, DataType.INT8 -> {
                val buffer = ShortBuffer.allocate(size)
                buffer.order()
                output[index] = buffer
            }
            DataType.INT64 -> {
                val buffer = LongBuffer.allocate(size)
                buffer.order()
                output[index] = buffer
            }
            DataType.BOOL, DataType.STRING -> {
                val buffer = ByteBuffer.allocateDirect(size * inputDataType.byteSize())
                buffer.order(ByteOrder.nativeOrder())
                output[index] = buffer
            }
        }
    }

But in this case, you cannot easily map ByteBuffer back to Array shape. This will fail:

 val out4 = output[3] as IntArray
 val out1 = output[4] as Array<Array<FloatArray>>
 val out2 = output[1] as Array<FloatArray>
 val out3 = output[5] as Array<FloatArray>

Question: Is it possible to convert ByteBuffer into a Multidimensional array?