kscience/kmath
9
0
Fork 1
Code Issues Pull Requests 6 Packages Projects Releases 16 Wiki Activity

Clone some APIs from DoubleTensorAlgebra to Nd4j, fix #348 #353

Open
CommanderTvis wants to merge 1 commits from commandertvis/nd4j into dev
pull from: commandertvis/nd4j
merge into: kscience:dev
Conversation 1 Commits 1 Files Changed 3 +186 -7
3 changed files with 186 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
kmath-nd4j/src/main/kotlin/space/kscience/kmath/nd4j/Nd4jArrayIterator.kt
View File

@ -15,9 +15,9 @@ private class Nd4jArrayIndicesIterator(private val iterateOver: INDArray) : Iter
override fun next(): IntArray {
val la = if (iterateOver.ordering() == 'c')
Shape.ind2subC(iterateOver, i++.toLong())!!
Shape.ind2subC(iterateOver, i++.toLong())
else
Shape.ind2sub(iterateOver, i++.toLong())!!
Shape.ind2sub(iterateOver, i++.toLong())
return la.toIntArray()
}

2
kmath-nd4j/src/main/kotlin/space/kscience/kmath/nd4j/Nd4jArrayStructure.kt
View File

@ -17,7 +17,7 @@ import space.kscience.kmath.nd.StructureND
*/
public sealed class Nd4jArrayStructure<T> : MutableStructureND<T> {
/**
* The wrapped [INDArray]. Since KMath uses [Int] indexes, assuming the size of [INDArray] is less or equal to
* The wrapped [INDArray]. Since KMath uses [Int] indices, assuming the size of [INDArray] is less or equal to
* [Int.MAX_VALUE].
*/
public abstract val ndArray: INDArray

187
kmath-nd4j/src/main/kotlin/space/kscience/kmath/nd4j/Nd4jTensorAlgebra.kt
View File

@ -9,6 +9,7 @@ import org.nd4j.linalg.api.ndarray.INDArray
import org.nd4j.linalg.api.ops.impl.summarystats.Variance
import org.nd4j.linalg.api.ops.impl.transforms.strict.ACosh
import org.nd4j.linalg.api.ops.impl.transforms.strict.ASinh
import org.nd4j.linalg.api.shape.Shape
import org.nd4j.linalg.factory.Nd4j
import org.nd4j.linalg.factory.ops.NDBase
import org.nd4j.linalg.ops.transforms.Transforms
@ -18,10 +19,14 @@ import space.kscience.kmath.nd.Shape
import space.kscience.kmath.nd.StructureND
import space.kscience.kmath.operations.DoubleField
import space.kscience.kmath.operations.Field
import space.kscience.kmath.samplers.GaussianSampler
import space.kscience.kmath.stat.RandomGenerator
import space.kscience.kmath.structures.toDoubleArray
import space.kscience.kmath.tensors.api.AnalyticTensorAlgebra
import space.kscience.kmath.tensors.api.Tensor
import space.kscience.kmath.tensors.api.TensorAlgebra
import space.kscience.kmath.tensors.core.DoubleTensorAlgebra
import space.kscience.kmath.tensors.core.DoubleTensor
import kotlin.math.abs
/**
* ND4J based [TensorAlgebra] implementation.
@ -156,10 +161,23 @@ public sealed interface Nd4jTensorAlgebra<T : Number, A : Field<T>> : AnalyticTe
Nd4j.getExecutioner().exec(Variance(ndArray, true, true, dim)).wrap()
private companion object {
private val ndBase: ThreadLocal<NDBase> = ThreadLocal.withInitial(::NDBase)
private val ndBase = NDBase()
}
private fun minusIndexFrom(n: Int, i: Int): Int = if (i >= 0) i else {
val ii = n + i
check(ii >= 0) { "Out of bound index $i for tensor of dim $n" }
ii
}
private fun getRandomNormals(n: Int, seed: Long): DoubleArray {
val distribution = GaussianSampler(0.0, 1.0)
val generator = RandomGenerator.default(seed)
return distribution.sample(generator).nextBufferBlocking(n).toDoubleArray()
}
/**
* [Double] specialization of [Nd4jTensorAlgebra].
*/
@ -180,7 +198,7 @@ public object DoubleNd4jTensorAlgebra : Nd4jTensorAlgebra<Double, DoubleField> {
@OptIn(PerformancePitfall::class)
override val StructureND<Double>.ndArray: INDArray
public override val StructureND<Double>.ndArray: INDArray
get() = when (this) {
is Nd4jArrayStructure<Double> -> ndArray
else -> Nd4j.zeros(*shape).also {
@ -197,7 +215,47 @@ public object DoubleNd4jTensorAlgebra : Nd4jTensorAlgebra<Double, DoubleField> {
offset: Int,
dim1: Int,
dim2: Int,
): Tensor<Double> = DoubleTensorAlgebra.diagonalEmbedding(diagonalEntries, offset, dim1, dim2)
): Tensor<Double> {
val diagonalEntriesNDArray = diagonalEntries.ndArray
val n = diagonalEntries.shape.size
val d1 = minusIndexFrom(n + 1, dim1)
val d2 = minusIndexFrom(n + 1, dim2)
check(d1 != d2) { "Diagonal dimensions cannot be identical $d1, $d2" }
check(d1 <= n && d2 <= n) { "Dimension out of range" }
var lessDim = d1
var greaterDim = d2
var realOffset = offset
if (lessDim > greaterDim) {
realOffset *= -1
lessDim = greaterDim.also { greaterDim = lessDim }
}
val resShape = diagonalEntries.shape.sliceArray(0 until lessDim) +
intArrayOf(diagonalEntries.shape[n - 1] + abs(realOffset)) +
diagonalEntries.shape.sliceArray(lessDim until greaterDim - 1) +
intArrayOf(diagonalEntries.shape[n - 1] + abs(realOffset)) +
diagonalEntries.shape.sliceArray(greaterDim - 1 until n - 1)
val resTensor = Nd4j.zeros(*resShape).wrap()
for (i in 0 until diagonalEntriesNDArray.length()) {
val multiIndex = (if (diagonalEntriesNDArray.ordering() == 'c')
Shape.ind2subC(diagonalEntriesNDArray, i)
else
Shape.ind2sub(diagonalEntriesNDArray, i)).toIntArray()
var offset1 = 0
var offset2 = abs(realOffset)
if (realOffset < 0) offset1 = offset2.also { offset2 = offset1 }
val diagonalMultiIndex = multiIndex.sliceArray(0 until lessDim) +
intArrayOf(multiIndex[n - 1] + offset1) +
multiIndex.sliceArray(lessDim until greaterDim - 1) +
intArrayOf(multiIndex[n - 1] + offset2) +
multiIndex.sliceArray(greaterDim - 1 until n - 1)
resTensor[diagonalMultiIndex] = diagonalEntries[multiIndex]
}
override fun StructureND<Double>.sum(): Double = ndArray.sumNumber().toDouble()
override fun StructureND<Double>.min(): Double = ndArray.minNumber().toDouble()
@ -205,4 +263,125 @@ public object DoubleNd4jTensorAlgebra : Nd4jTensorAlgebra<Double, DoubleField> {
override fun StructureND<Double>.mean(): Double = ndArray.meanNumber().toDouble()
override fun StructureND<Double>.std(): Double = ndArray.stdNumber().toDouble()
override fun StructureND<Double>.variance(): Double = ndArray.varNumber().toDouble()
return resTensor
}
/**
* Compares element-wise two tensors with a specified precision.
*
* @param other the tensor to compare with `input` tensor.
* @param epsilon permissible error when comparing two Double values.
* @return true if two tensors have the same shape and elements, false otherwise.
*/
public fun Tensor<Double>.eq(other: Tensor<Double>, epsilon: Double): Boolean =
ndArray.equalsWithEps(other, epsilon)
/**
* Compares element-wise two tensors.
* Comparison of two Double values occurs with 1e-5 precision.
*
* @param other the tensor to compare with `input` tensor.
* @return true if two tensors have the same shape and elements, false otherwise.
*/
public infix fun Tensor<Double>.eq(other: Tensor<Double>): Boolean = eq(other, 1e-5)
public override fun Tensor<Double>.sum(): Double = ndArray.sumNumber().toDouble()
public override fun Tensor<Double>.min(): Double = ndArray.minNumber().toDouble()
public override fun Tensor<Double>.max(): Double = ndArray.maxNumber().toDouble()
public override fun Tensor<Double>.mean(): Double = ndArray.meanNumber().toDouble()
public override fun Tensor<Double>.std(): Double = ndArray.stdNumber().toDouble()
public override fun Tensor<Double>.variance(): Double = ndArray.varNumber().toDouble()
/**
* Constructs a tensor with the specified shape and data.
*
* @param shape the desired shape for the tensor.
* @param buffer one-dimensional data array.
* @return tensor with the [shape] shape and [buffer] data.
*/
public fun fromArray(shape: IntArray, buffer: DoubleArray): Nd4jArrayStructure<Double> =
Nd4j.create(buffer, shape).wrap()
/**
* Constructs a tensor with the specified shape and initializer.
*
* @param shape the desired shape for the tensor.
* @param initializer mapping tensor indices to values.
* @return tensor with the [shape] shape and data generated by the [initializer].
*/
public fun produce(shape: IntArray, initializer: (IntArray) -> Double): Nd4jArrayStructure<Double> {
val struct = Nd4j.create(*shape)!!.wrap()
struct.indicesIterator().forEach { struct[it] = initializer(it) }
return struct
}
/**
* Returns a tensor of random numbers drawn from normal distributions with 0.0 mean and 1.0 standard deviation.
*
* @param shape the desired shape for the output tensor.
* @param seed the random seed of the pseudo-random number generator.
* @return tensor of a given shape filled with numbers from the normal distribution
* with 0.0 mean and 1.0 standard deviation.
*/
public fun randomNormal(shape: IntArray, seed: Long = 0): Nd4jArrayStructure<Double> =
fromArray(shape, getRandomNormals(shape.reduce(Int::times), seed))
/**
* Returns a tensor with the same shape as `input` of random numbers drawn from normal distributions
* with 0.0 mean and 1.0 standard deviation.
*
* @param seed the random seed of the pseudo-random number generator.
* @return tensor with the same shape as `input` filled with numbers from the normal distribution
* with 0.0 mean and 1.0 standard deviation.
*/
public fun Tensor<Double>.randomNormalLike(seed: Long = 0): Nd4jArrayStructure<Double> =
fromArray(shape, getRandomNormals(shape.reduce(Int::times), seed))
/**
* Creates a tensor of a given shape and fills all elements with a given value.
*
* @param value the value to fill the output tensor with.
* @param shape array of integers defining the shape of the output tensor.
* @return tensor with the [shape] shape and filled with [value].
*/
public fun full(value: Double, shape: IntArray): Nd4jArrayStructure<Double> = Nd4j.valueArrayOf(shape, value).wrap()
/**
* Returns a tensor with the same shape as `input` filled with [value].
*
* @param value the value to fill the output tensor with.
* @return tensor with the `input` tensor shape and filled with [value].
*/
public fun Tensor<Double>.fullLike(value: Double): Nd4jArrayStructure<Double> =
Nd4j.valueArrayOf(ndArray.shape(), value).wrap()
/**
* Returns a tensor filled with the scalar value 0.0, with the shape defined by the variable argument [shape].
*
* @param shape array of integers defining the shape of the output tensor.
* @return tensor filled with the scalar value 0.0, with the [shape] shape.
*/
public fun zeros(shape: IntArray): Nd4jArrayStructure<Double> = full(0.0, shape)
/**
* Returns a tensor filled with the scalar value 0.0, with the same shape as a given array.
*
* @return tensor filled with the scalar value 0.0, with the same shape as `input` tensor.
*/
public fun Tensor<Double>.zeroesLike(): Nd4jArrayStructure<Double> = Nd4j.zerosLike(ndArray).wrap()
/**
* Returns a tensor filled with the scalar value 1.0, with the shape defined by the variable argument [shape].
*
* @param shape array of integers defining the shape of the output tensor.
* @return tensor filled with the scalar value 1.0, with the [shape] shape.
*/
public fun ones(shape: IntArray): Nd4jArrayStructure<Double> = Nd4j.ones(*shape).wrap()
/**
* Returns a tensor filled with the scalar value 1.0, with the same shape as a given array.
*
* @return tensor filled with the scalar value 1.0, with the same shape as `input` tensor.
*/
public fun Tensor<Double>.onesLike(): Nd4jArrayStructure<Double> = Nd4j.onesLike(ndArray).wrap()
}