This commit is contained in:
Andrei Kislitsyn 2021-05-07 15:53:09 +03:00
parent d73d03c055
commit 5b725a087b

View File

@ -35,3 +35,35 @@ public fun IntTensor.toIntArray(): IntArray {
mutableBuffer[bufferStart + i] mutableBuffer[bufferStart + i]
} }
} }
/**
* Casts [Array<DoubleArray>] to [DoubleTensor]
*/
public fun Array<DoubleArray>.toDoubleTensor(): DoubleTensor {
val n = size
check(n > 0) { "An empty array cannot be casted to tensor" }
val m = first().size
check(m > 0) { "Inner arrays must have at least 1 argument" }
check(all { size == m }) { "Inner arrays must be the same size" }
val shape = intArrayOf(n, m)
val buffer = this.flatMap { arr -> arr.map { it } }.toDoubleArray()
return DoubleTensor(shape, buffer, 0)
}
/**
* Casts [Array<IntArray>] to [IntTensor]
*/
public fun Array<IntArray>.toIntTensor(): IntTensor {
val n = size
check(n > 0) { "An empty array cannot be casted to tensor" }
val m = first().size
check(m > 0) { "Inner arrays must have at least 1 argument" }
check(all { size == m }) { "Inner arrays must be the same size" }
val shape = intArrayOf(n, m)
val buffer = this.flatMap { arr -> arr.map { it } }.toIntArray()
return IntTensor(shape, buffer, 0)
}