KMP library for tensors #300
@ -5,14 +5,107 @@
|
|||||||
|
|
||||||
package space.kscience.kmath.tensors.api
|
package space.kscience.kmath.tensors.api
|
||||||
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Element-wise analytic operations on [Tensor].
|
* Analytic operations on [Tensor].
|
||||||
*
|
*
|
||||||
* @param T the type of items closed under analytic functions in the tensors.
|
* @param T the type of items closed under analytic functions in the tensors.
|
||||||
*/
|
*/
|
||||||
public interface AnalyticTensorAlgebra<T> :
|
public interface AnalyticTensorAlgebra<T> :
|
||||||
TensorPartialDivisionAlgebra<T> {
|
TensorPartialDivisionAlgebra<T> {
|
||||||
|
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the minimum value of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.min(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the minimum value of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the minimum value of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.min(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the maximum value of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.max(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the maximum value of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the maximum value of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.max(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the mean of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.mean(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the mean of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the mean of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.mean(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the standard deviation of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.std(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the standard deviation of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the standard deviation of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.std(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the variance of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.variance(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the variance of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the variance of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.variance(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
//For information: https://pytorch.org/docs/stable/generated/torch.exp.html
|
//For information: https://pytorch.org/docs/stable/generated/torch.exp.html
|
||||||
public fun Tensor<T>.exp(): Tensor<T>
|
public fun Tensor<T>.exp(): Tensor<T>
|
||||||
|
|
||||||
|
@ -1,124 +0,0 @@
|
|||||||
/*
|
|
||||||
* Copyright 2018-2021 KMath contributors.
|
|
||||||
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
|
|
||||||
*/
|
|
||||||
|
|
||||||
package space.kscience.kmath.tensors.api
|
|
||||||
|
|
||||||
import space.kscience.kmath.tensors.core.DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Common algebra with statistics methods. Operates on [Tensor].
|
|
||||||
*/
|
|
||||||
|
|
||||||
public interface StatisticTensorAlgebra<T>: TensorAlgebra<T> {
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the minimum value of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.min(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the minimum value of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the minimum value of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.min(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the maximum value of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.max(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the maximum value of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the maximum value of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.max(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the sum of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.sum(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the sum of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the sum of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.sum(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the mean of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.mean(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the mean of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the mean of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.mean(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the standard deviation of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.std(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the standard deviation of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the standard deviation of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.std(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the variance of all elements in the input tensor.
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.variance(): Double
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Returns the variance of each row of the input tensor in the given dimension [dim].
|
|
||||||
*
|
|
||||||
* If [keepDim] is true, the output tensor is of the same size as
|
|
||||||
* input except in the dimension [dim] where it is of size 1.
|
|
||||||
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
|
||||||
*
|
|
||||||
* @param dim the dimension to reduce.
|
|
||||||
* @param keepDim whether the output tensor has [dim] retained or not.
|
|
||||||
* @return the variance of each row of the input tensor in the given dimension [dim].
|
|
||||||
*/
|
|
||||||
public fun Tensor<T>.variance(dim: Int, keepDim: Boolean): DoubleTensor
|
|
||||||
|
|
||||||
}
|
|
@ -251,4 +251,24 @@ public interface TensorAlgebra<T>: Algebra<Tensor<T>> {
|
|||||||
dim2: Int = -1
|
dim2: Int = -1
|
||||||
): Tensor<T>
|
): Tensor<T>
|
||||||
|
|
||||||
|
/**
|
||||||
|
* @return the sum of all elements in the input tensor.
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.sum(): T
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Returns the sum of each row of the input tensor in the given dimension [dim].
|
||||||
|
*
|
||||||
|
* If [keepDim] is true, the output tensor is of the same size as
|
||||||
|
* input except in the dimension [dim] where it is of size 1.
|
||||||
|
* Otherwise, [dim] is squeezed, resulting in the output tensor having 1 fewer dimension.
|
||||||
|
*
|
||||||
|
* @param dim the dimension to reduce.
|
||||||
|
* @param keepDim whether the output tensor has [dim] retained or not.
|
||||||
|
* @return the sum of each row of the input tensor in the given dimension [dim].
|
||||||
|
*/
|
||||||
|
public fun Tensor<T>.sum(dim: Int, keepDim: Boolean): Tensor<T>
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
}
|
}
|
||||||
|
@ -14,6 +14,60 @@ import kotlin.math.*
|
|||||||
public object DoubleAnalyticTensorAlgebra :
|
public object DoubleAnalyticTensorAlgebra :
|
||||||
AnalyticTensorAlgebra<Double>,
|
AnalyticTensorAlgebra<Double>,
|
||||||
DoubleTensorAlgebra() {
|
DoubleTensorAlgebra() {
|
||||||
|
|
||||||
|
override fun Tensor<Double>.min(): Double = this.fold { it.minOrNull()!! }
|
||||||
|
|
||||||
|
override fun Tensor<Double>.min(dim: Int, keepDim: Boolean): DoubleTensor =
|
||||||
|
foldDim({ x -> x.minOrNull()!! }, dim, keepDim)
|
||||||
|
|
||||||
|
override fun Tensor<Double>.max(): Double = this.fold { it.maxOrNull()!! }
|
||||||
|
|
||||||
|
override fun Tensor<Double>.max(dim: Int, keepDim: Boolean): DoubleTensor =
|
||||||
|
foldDim({ x -> x.maxOrNull()!! }, dim, keepDim)
|
||||||
|
|
||||||
|
|
||||||
|
override fun Tensor<Double>.mean(): Double = this.fold { it.sum() / tensor.numElements }
|
||||||
|
|
||||||
|
override fun Tensor<Double>.mean(dim: Int, keepDim: Boolean): DoubleTensor =
|
||||||
|
foldDim(
|
||||||
|
{ arr ->
|
||||||
|
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
||||||
|
arr.sum() / shape[dim]
|
||||||
|
},
|
||||||
|
dim,
|
||||||
|
keepDim
|
||||||
|
)
|
||||||
|
|
||||||
|
override fun Tensor<Double>.std(): Double = this.fold { arr ->
|
||||||
|
val mean = arr.sum() / tensor.numElements
|
||||||
|
sqrt(arr.sumOf { (it - mean) * (it - mean) } / (tensor.numElements - 1))
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun Tensor<Double>.std(dim: Int, keepDim: Boolean): DoubleTensor = foldDim(
|
||||||
|
{ arr ->
|
||||||
|
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
||||||
|
val mean = arr.sum() / shape[dim]
|
||||||
|
sqrt(arr.sumOf { (it - mean) * (it - mean) } / (shape[dim] - 1))
|
||||||
|
},
|
||||||
|
dim,
|
||||||
|
keepDim
|
||||||
|
)
|
||||||
|
|
||||||
|
override fun Tensor<Double>.variance(): Double = this.fold { arr ->
|
||||||
|
val mean = arr.sum() / tensor.numElements
|
||||||
|
arr.sumOf { (it - mean) * (it - mean) } / (tensor.numElements - 1)
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun Tensor<Double>.variance(dim: Int, keepDim: Boolean): DoubleTensor = foldDim(
|
||||||
|
{ arr ->
|
||||||
|
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
||||||
|
val mean = arr.sum() / shape[dim]
|
||||||
|
arr.sumOf { (it - mean) * (it - mean) } / (shape[dim] - 1)
|
||||||
|
},
|
||||||
|
dim,
|
||||||
|
keepDim
|
||||||
|
)
|
||||||
|
|
||||||
override fun Tensor<Double>.exp(): DoubleTensor = tensor.map(::exp)
|
override fun Tensor<Double>.exp(): DoubleTensor = tensor.map(::exp)
|
||||||
|
|
||||||
override fun Tensor<Double>.log(): DoubleTensor = tensor.map(::ln)
|
override fun Tensor<Double>.log(): DoubleTensor = tensor.map(::ln)
|
||||||
|
@ -1,104 +0,0 @@
|
|||||||
/*
|
|
||||||
* Copyright 2018-2021 KMath contributors.
|
|
||||||
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
|
|
||||||
*/
|
|
||||||
|
|
||||||
package space.kscience.kmath.tensors.core.algebras
|
|
||||||
|
|
||||||
import kotlin.math.sqrt
|
|
||||||
|
|
||||||
import space.kscience.kmath.tensors.api.*
|
|
||||||
import space.kscience.kmath.tensors.core.*
|
|
||||||
import space.kscience.kmath.tensors.core.algebras.DoubleStatisticTensorAlgebra.max
|
|
||||||
import space.kscience.kmath.tensors.core.algebras.DoubleStatisticTensorAlgebra.mean
|
|
||||||
import space.kscience.kmath.tensors.core.algebras.DoubleStatisticTensorAlgebra.min
|
|
||||||
import space.kscience.kmath.tensors.core.algebras.DoubleStatisticTensorAlgebra.sum
|
|
||||||
import space.kscience.kmath.tensors.core.algebras.DoubleStatisticTensorAlgebra.variance
|
|
||||||
|
|
||||||
public object DoubleStatisticTensorAlgebra : StatisticTensorAlgebra<Double>, DoubleTensorAlgebra() {
|
|
||||||
|
|
||||||
private fun Tensor<Double>.fold(foldFunction: (DoubleArray) -> Double): Double =
|
|
||||||
foldFunction(this.tensor.toDoubleArray())
|
|
||||||
|
|
||||||
private fun Tensor<Double>.foldDim(
|
|
||||||
foldFunction: (DoubleArray) -> Double,
|
|
||||||
dim: Int,
|
|
||||||
keepDim: Boolean
|
|
||||||
): DoubleTensor {
|
|
||||||
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
|
||||||
val resShape = if (keepDim) {
|
|
||||||
shape.take(dim).toIntArray() + intArrayOf(1) + shape.takeLast(dimension - dim - 1).toIntArray()
|
|
||||||
} else {
|
|
||||||
shape.take(dim).toIntArray() + shape.takeLast(dimension - dim - 1).toIntArray()
|
|
||||||
}
|
|
||||||
val resNumElements = resShape.reduce(Int::times)
|
|
||||||
val resTensor = DoubleTensor(resShape, DoubleArray(resNumElements) { 0.0 }, 0)
|
|
||||||
for (index in resTensor.linearStructure.indices()) {
|
|
||||||
val prefix = index.take(dim).toIntArray()
|
|
||||||
val suffix = index.takeLast(dimension - dim - 1).toIntArray()
|
|
||||||
resTensor[index] = foldFunction(DoubleArray(shape[dim]) { i ->
|
|
||||||
this[prefix + intArrayOf(i) + suffix]
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
return resTensor
|
|
||||||
}
|
|
||||||
|
|
||||||
override fun Tensor<Double>.min(): Double = this.fold { it.minOrNull()!! }
|
|
||||||
|
|
||||||
override fun Tensor<Double>.min(dim: Int, keepDim: Boolean): DoubleTensor =
|
|
||||||
foldDim({ x -> x.minOrNull()!! }, dim, keepDim)
|
|
||||||
|
|
||||||
override fun Tensor<Double>.max(): Double = this.fold { it.maxOrNull()!! }
|
|
||||||
|
|
||||||
override fun Tensor<Double>.max(dim: Int, keepDim: Boolean): DoubleTensor =
|
|
||||||
foldDim({ x -> x.maxOrNull()!! }, dim, keepDim)
|
|
||||||
|
|
||||||
override fun Tensor<Double>.sum(): Double = this.fold { it.sum() }
|
|
||||||
|
|
||||||
override fun Tensor<Double>.sum(dim: Int, keepDim: Boolean): DoubleTensor =
|
|
||||||
foldDim({ x -> x.sum() }, dim, keepDim)
|
|
||||||
|
|
||||||
override fun Tensor<Double>.mean(): Double = this.fold { it.sum() / tensor.numElements }
|
|
||||||
|
|
||||||
override fun Tensor<Double>.mean(dim: Int, keepDim: Boolean): DoubleTensor =
|
|
||||||
foldDim(
|
|
||||||
{ arr ->
|
|
||||||
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
|
||||||
arr.sum() / shape[dim]
|
|
||||||
},
|
|
||||||
dim,
|
|
||||||
keepDim
|
|
||||||
)
|
|
||||||
|
|
||||||
override fun Tensor<Double>.std(): Double = this.fold { arr ->
|
|
||||||
val mean = arr.sum() / tensor.numElements
|
|
||||||
sqrt(arr.sumOf { (it - mean) * (it - mean) } / (tensor.numElements - 1))
|
|
||||||
}
|
|
||||||
|
|
||||||
override fun Tensor<Double>.std(dim: Int, keepDim: Boolean): DoubleTensor = foldDim(
|
|
||||||
{ arr ->
|
|
||||||
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
|
||||||
val mean = arr.sum() / shape[dim]
|
|
||||||
sqrt(arr.sumOf { (it - mean) * (it - mean) } / (shape[dim] - 1))
|
|
||||||
},
|
|
||||||
dim,
|
|
||||||
keepDim
|
|
||||||
)
|
|
||||||
|
|
||||||
override fun Tensor<Double>.variance(): Double = this.fold { arr ->
|
|
||||||
val mean = arr.sum() / tensor.numElements
|
|
||||||
arr.sumOf { (it - mean) * (it - mean) } / (tensor.numElements - 1)
|
|
||||||
}
|
|
||||||
|
|
||||||
override fun Tensor<Double>.variance(dim: Int, keepDim: Boolean): DoubleTensor = foldDim(
|
|
||||||
{ arr ->
|
|
||||||
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
|
||||||
val mean = arr.sum() / shape[dim]
|
|
||||||
arr.sumOf { (it - mean) * (it - mean) } / (shape[dim] - 1)
|
|
||||||
},
|
|
||||||
dim,
|
|
||||||
keepDim
|
|
||||||
)
|
|
||||||
|
|
||||||
}
|
|
@ -415,4 +415,36 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra<Double> {
|
|||||||
public fun Tensor<Double>.rowsByIndices(indices: IntArray): DoubleTensor {
|
public fun Tensor<Double>.rowsByIndices(indices: IntArray): DoubleTensor {
|
||||||
return stack(indices.map { this[it] })
|
return stack(indices.map { this[it] })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
internal fun Tensor<Double>.fold(foldFunction: (DoubleArray) -> Double): Double =
|
||||||
|
foldFunction(tensor.toDoubleArray())
|
||||||
|
|
||||||
|
internal fun Tensor<Double>.foldDim(
|
||||||
|
foldFunction: (DoubleArray) -> Double,
|
||||||
|
dim: Int,
|
||||||
|
keepDim: Boolean
|
||||||
|
): DoubleTensor {
|
||||||
|
check(dim < dimension) { "Dimension $dim out of range $dimension" }
|
||||||
|
val resShape = if (keepDim) {
|
||||||
|
shape.take(dim).toIntArray() + intArrayOf(1) + shape.takeLast(dimension - dim - 1).toIntArray()
|
||||||
|
} else {
|
||||||
|
shape.take(dim).toIntArray() + shape.takeLast(dimension - dim - 1).toIntArray()
|
||||||
|
}
|
||||||
|
val resNumElements = resShape.reduce(Int::times)
|
||||||
|
val resTensor = DoubleTensor(resShape, DoubleArray(resNumElements) { 0.0 }, 0)
|
||||||
|
for (index in resTensor.linearStructure.indices()) {
|
||||||
|
val prefix = index.take(dim).toIntArray()
|
||||||
|
val suffix = index.takeLast(dimension - dim - 1).toIntArray()
|
||||||
|
resTensor[index] = foldFunction(DoubleArray(shape[dim]) { i ->
|
||||||
|
tensor[prefix + intArrayOf(i) + suffix]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return resTensor
|
||||||
|
}
|
||||||
|
|
||||||
|
override fun Tensor<Double>.sum(): Double = tensor.fold { it.sum() }
|
||||||
|
|
||||||
|
override fun Tensor<Double>.sum(dim: Int, keepDim: Boolean): DoubleTensor =
|
||||||
|
foldDim({ x -> x.sum() }, dim, keepDim)
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user