diff --git a/README.md b/README.md index 89fc6cb8..4cf6393e 100644 --- a/README.md +++ b/README.md @@ -34,9 +34,9 @@ trait Batch derives Label trait Feature derives Label // Create a 2D tensor with shape (3, 2), labeled with Batch and Feature -val t = Tensor.fromArray( +val t = Tensor( Shape(Axis[Batch] -> 3, Axis[Feature] -> 2), - VType[Float])( +).fromArray( Array( 1.0f, 2.0f, 3.0f, 4.0f, diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index 89dff7a5..027e8a66 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -61,7 +61,7 @@ package object dimwit: export dimwit.tensor.{Tensor, Tensor0, Tensor1, Tensor2, Tensor3} export dimwit.tensor.{Shape, Shape0, Shape1, Shape2, Shape3} export dimwit.tensor.{DType, Device} - export dimwit.tensor.{VType, ExecutionType, Label, Labels, Axis, AxisIndex, AxisIndices, Dim} + export dimwit.tensor.{VType, ExecutionType, ExecutionTypeFor, Label, Labels, Axis, AxisIndex, AxisIndices, Dim} // Export operations export dimwit.tensor.TensorOps.* diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index c5c87e23..cd3968b2 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -23,9 +23,11 @@ class Normal[LocT <: T, ScaleT <: T, T <: Tuple: Labels]( standardNormal * scale + loc object Normal: - def standardNormal[T <: Tuple: Labels](shape: Shape[T]) = new Normal( - loc = Tensor.zeros(shape, VType[Float]), - scale = Tensor.ones(shape, VType[Float]) + + def standardSample(key: Random.Key): Tensor0[Float] = standardNormal(Shape.empty).sample(key) + def standardNormal[T <: Tuple: Labels](shape: Shape[T])(using executionType: ExecutionType[Float]) = new Normal( + loc = Tensor(shape).fill(0f), + scale = Tensor(shape).fill(1f) ) class Uniform[T <: Tuple: Labels]( diff --git a/core/src/main/scala/dimwit/tensor/ArrayWriter.scala b/core/src/main/scala/dimwit/tensor/ArrayWriter.scala new file mode 100644 index 00000000..8ff22e22 --- /dev/null +++ b/core/src/main/scala/dimwit/tensor/ArrayWriter.scala @@ -0,0 +1,93 @@ +package dimwit.tensor + +import java.nio.ByteBuffer +import java.util.Base64 +import java.nio.ByteOrder +import me.shadaj.scalapy.py +import me.shadaj.scalapy.py.SeqConverters +import me.shadaj.scalapy.readwrite.Writer +import me.shadaj.scalapy.interpreter.PyValue +import dimwit.jax.Jax + +trait WriterEvidence[A]: + type V + +/** Type class for providing evidence that a scalar of type A can be converted to type V using a ScalaPy Writer in a Tensor context. + * The type A is the input scalar type, allowing to define an internal precision (dtype) based on the scalar type. + * For example creating a Tensor[?, Int] from a scalar of type Byte with internal dtype uint8, int8, int16 or int32 (based on given ExecutionType[Byte]). + * + * The type V is the value type of the resulting Tensor (should be Boolean, Int or Float; or custom opaque types). + */ +object WriterEvidence: + + type Aux[A, V0] = WriterEvidence[A] { type V = V0 } + +// Helper to instantiate + def apply[A, V0]: Aux[A, V0] = new WriterEvidence[A]: + type V = V0 + + given Aux[Float, Float] = apply + given Aux[Int, Int] = apply + given Aux[Boolean, Boolean] = apply + given Aux[Double, Float] = apply // Double casts to Float + given Aux[Byte, Int] = apply // Byte casts to Int + +/** Type class for creating Tensors of different value types from arrays of different base types. + * While allowing to define an internal precision (dtype) based on the array type. + * For example creating a Tensor[?, Int] from an Array[Byte] with internal dtype uint8, int8, int16 or int32 (based on given ExecutionType[Byte]). + * + * @param A The base type of the input array. + * @param V The value type of the resulting Tensor. + */ +trait ArrayWriter[A]: + type V + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[A]): Tensor[T, V] + +object ArrayWriter: + type Aux[A, V0] = ArrayWriter[A] { type V = V0 } + + val base64Loader = py.eval("lambda b64, shape, dtype: __import__('jax').numpy.array(__import__('numpy').frombuffer(__import__('base64').b64decode(b64), dtype=dtype).reshape(shape))") + + private def byteArrayToTensor[T <: Tuple: Labels, V](shape: Shape[T], byteArray: Array[Byte], jaxDType: Jax.PyDynamic): Tensor[T, V] = + val b64String = Base64.getEncoder.encodeToString(byteArray) + Tensor(base64Loader(b64String, shape.dimensions.toPythonProxy, jaxDType)) + + given (using ExecutionType[Double]): ArrayWriter.Aux[Double, Float] = new ArrayWriter[Double]: + type V = Float + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Double]): Tensor[T, Float] = + require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") + val dtype = ExecutionType[Double].dtype + val byteArray = dtype.write(values) + byteArrayToTensor(shape, byteArray, dtype.jaxType) + + given (using ExecutionType[Float]): ArrayWriter.Aux[Float, Float] = new ArrayWriter[Float]: + type V = Float + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Float]): Tensor[T, Float] = + require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") + val dtype = ExecutionType[Float].dtype + val byteArray = dtype.write(values) + byteArrayToTensor(shape, byteArray, dtype.jaxType) + + given (using ExecutionType[Int]): ArrayWriter.Aux[Int, Int] = new ArrayWriter[Int]: + type V = Int + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Int]): Tensor[T, Int] = + require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") + val dtype = ExecutionType[Int].dtype + val byteArray = dtype.write(values) + byteArrayToTensor(shape, byteArray, dtype.jaxType) + + given (using ExecutionType[Byte]): ArrayWriter.Aux[Byte, Int] = new ArrayWriter[Byte]: + type V = Int + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Byte]): Tensor[T, Int] = + require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") + val dtype = ExecutionType[Byte].dtype + val byteArray = dtype.write(values) + byteArrayToTensor(shape, byteArray, dtype.jaxType) + + given (using ExecutionType[Boolean]): ArrayWriter.Aux[Boolean, Boolean] = new ArrayWriter[Boolean]: + type V = Boolean + def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Boolean]): Tensor[T, Boolean] = + require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") + val dtype = ExecutionType[Boolean].dtype + val byteArray = dtype.write(values) + byteArrayToTensor(shape, byteArray, dtype.jaxType) diff --git a/core/src/main/scala/dimwit/tensor/DType.scala b/core/src/main/scala/dimwit/tensor/DType.scala index 749b8223..4e1ccbae 100644 --- a/core/src/main/scala/dimwit/tensor/DType.scala +++ b/core/src/main/scala/dimwit/tensor/DType.scala @@ -1,5 +1,7 @@ package dimwit.tensor import dimwit.jax.JaxDType +import java.nio.ByteBuffer +import java.nio.ByteOrder enum DType(val name: String, val size: Int): case Float32 extends DType("float32", 4) @@ -16,3 +18,59 @@ enum DType(val name: String, val size: Int): case Complex128 extends DType("complex128", 16) lazy val jaxType = JaxDType.jaxDtype(this) + + /** Writes the scala array into the buffer according to THIS DType's format. + * Handles conversions (e.g. Array[Double] -> Float32 buffer). + */ + def write(values: Array[?]): Array[Byte] = + val buffer = ByteBuffer.allocate(values.length * this.size) + buffer.order(ByteOrder.LITTLE_ENDIAN) // Python uses little-endian + + // write values into buffer according to this DType + (this, values) match + // --- Float32 Target --- + case (Float32, arr: Array[Float]) => + buffer.asFloatBuffer().put(arr) + case (Float32, arr: Array[Double]) => + val fb = buffer.asFloatBuffer() + var i = 0 + while i < arr.length do + fb.put(arr(i).toFloat); i += 1 + + // --- Float64 Target --- + case (Float64, arr: Array[Double]) => + buffer.asDoubleBuffer().put(arr) + case (Float64, arr: Array[Float]) => + val db = buffer.asDoubleBuffer() + var i = 0; + while i < arr.length do + db.put(arr(i).toDouble); i += 1 + + // --- Int32 Target --- + case (Int32, arr: Array[Int]) => + buffer.asIntBuffer().put(arr) + case (Int32, arr: Array[Byte]) => + val ib = buffer.asIntBuffer() + var i = 0; + while i < arr.length do + ib.put(arr(i).toInt); i += 1 + + // --- Int8 Target --- + case (Int8, arr: Array[Byte]) => + buffer.put(arr) + + case (UInt8, arr: Array[Byte]) => + buffer.put(arr) // Interpret bytes as unsigned when reading + + // --- Bool Target (1 byte per bool) --- + case (Bool, arr: Array[Boolean]) => + var i = 0 + while i < arr.length do + buffer.put(if arr(i) then 1.toByte else 0.toByte) + i += 1 + + // --- Fallback/Error --- + case _ => + throw new IllegalArgumentException(s"Conversion from ${values.getClass.getSimpleName} to DType $name is not supported or implemented.") + + buffer.array() diff --git a/core/src/main/scala/dimwit/tensor/Tensor.scala b/core/src/main/scala/dimwit/tensor/Tensor.scala index 2e1dba6e..4711c9f5 100644 --- a/core/src/main/scala/dimwit/tensor/Tensor.scala +++ b/core/src/main/scala/dimwit/tensor/Tensor.scala @@ -13,9 +13,7 @@ import dimwit.stats.{Normal, Uniform} import me.shadaj.scalapy.readwrite.Writer import scala.reflect.ClassTag import scala.annotation.unchecked.uncheckedVariance -import java.nio.ByteBuffer -import java.util.Base64 -import java.nio.ByteOrder +import dimwit.stats.IndependentDistribution enum Device(val platform: String): case CPU extends Device("cpu") @@ -74,58 +72,28 @@ object Tensor: type IndicesOf[T <: Tuple] = Tuple.Map[T, [_] =>> Int] + case class Factory[T <: Tuple: Labels](val shape: Shape[T]): + + def fill[A: ExecutionType: Writer, V](value: A)(using ev: WriterEvidence.Aux[A, V]): Tensor[T, V] = + Tensor(Jax.jnp.full(shape.dimensions.toPythonProxy, value, dtype = ExecutionType[A].dtype.jaxType)) + + def fromArray[A: ExecutionType, V](values: Array[A])(using t2a: ArrayWriter.Aux[A, V]): Tensor[T, V] = + t2a.fromArray[T](shape)(values) + + case class LikeFactory[T <: Tuple: Labels, V](val other: Tensor[T, V]): + + def fill[A: Writer](value: A)(using ev: WriterEvidence.Aux[A, V]): Tensor[T, V] = + Tensor(Jax.jnp.full(other.shape.dimensions.toPythonProxy, value, dtype = other.dtype.jaxType)) + + def fromArray[A](values: Array[A])(using t2a: ArrayWriter.Aux[A, V]): Tensor[T, V] = + given ExecutionType[A] = ExecutionTypeFor[A](other.dtype) // fix the underlying dtype to match the other tensor's dtype + summon[ArrayWriter[A]].fromArray[T](other.shape)(values) + + def apply[T <: Tuple: Labels](shape: Shape[T]): Tensor.Factory[T] = Tensor.Factory(shape) def apply[T <: Tuple: Labels, V](jaxValue: Jax.PyDynamic): Tensor[T, V] = new Tensor(jaxValue) - def randn[T <: Tuple: Labels](shape: Shape[T])(key: Random.Key)(using - executionType: ExecutionType[Float] - ): Tensor[T, Float] = Normal.standardNormal(shape).sample(key) + def like[T <: Tuple: Labels, V](template: Tensor[T, V]): Tensor.LikeFactory[T, V] = Tensor.LikeFactory(template) def fromPy[T <: Tuple: Labels, V](vtype: VType[V])(jaxValue: Jax.PyDynamic): Tensor[T, V] = new Tensor(jaxValue) - def zeros[T <: Tuple: Labels, V](shape: Shape[T], vtype: VType[V]): Tensor[T, V] = Tensor(Jax.jnp.zeros(shape.dimensions.toPythonProxy, dtype = vtype.dtype.jaxType)) - def ones[T <: Tuple: Labels, V](shape: Shape[T], vtype: VType[V]): Tensor[T, V] = Tensor(Jax.jnp.ones(shape.dimensions.toPythonProxy, dtype = vtype.dtype.jaxType)) - def const[T <: Tuple: Labels, V](shape: Shape[T], vtype: VType[V])(value: V)(using writer: Writer[V]): Tensor[T, V] = Tensor(Jax.jnp.full(shape.dimensions.toPythonProxy, value, dtype = vtype.dtype.jaxType)) - - def fromArray[T <: Tuple: Labels](shape: Shape[T], vtype: VType[Float])(values: Array[Float]): Tensor[T, Float] = fromFloatArray(shape)(values) - def fromArray[T <: Tuple: Labels](shape: Shape[T], vtype: VType[Int])(values: Array[Int]): Tensor[T, Int] = fromIntArray(shape)(values) - def fromArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Byte]): Tensor[T, Int] = fromByteArray(shape)(values) - def fromArray[T <: Tuple: Labels](shape: Shape[T], vtype: VType[Boolean])(values: Array[Boolean]): Tensor[T, Boolean] = fromBooleanArray(shape)(values) - - /** array.toPythonProxy is very inefficient for large arrays, so we use base64 encoding as a workaround */ - private val base64Loader = py.eval("lambda b64, shape, dtype: __import__('jax').numpy.array(__import__('numpy').frombuffer(__import__('base64').b64decode(b64), dtype=dtype).reshape(shape))") - - def fromFloatArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Float]): Tensor[T, Float] = - require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") - val floatArr = values.asInstanceOf[Array[Float]] - val buffer = ByteBuffer.allocate(floatArr.length * 4) - buffer.order(ByteOrder.LITTLE_ENDIAN) - buffer.asFloatBuffer().put(floatArr) - val b64String = Base64.getEncoder.encodeToString(buffer.array()) - Tensor(base64Loader(b64String, shape.dimensions.toPythonProxy, "float32")) - - def fromIntArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Int]): Tensor[T, Int] = - require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") - val intArr = values.asInstanceOf[Array[Int]] - val buffer = ByteBuffer.allocate(intArr.length * 4) - buffer.order(ByteOrder.LITTLE_ENDIAN) - buffer.asIntBuffer().put(intArr) - val b64String = Base64.getEncoder.encodeToString(buffer.array()) - Tensor(base64Loader(b64String, shape.dimensions.toPythonProxy, "int32")) - - def fromByteArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Byte]): Tensor[T, Int] = - require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") - val buffer = ByteBuffer.allocate(values.length) - buffer.order(ByteOrder.LITTLE_ENDIAN) - buffer.put(values) - val b64String = Base64.getEncoder.encodeToString(buffer.array()) - Tensor(base64Loader(b64String, shape.dimensions.toPythonProxy, "uint8")) - - def fromBooleanArray[T <: Tuple: Labels](shape: Shape[T])(values: Array[Boolean]): Tensor[T, Boolean] = - require(values.length == shape.size, s"Values length ${values.length} does not match shape size ${shape.size}") - val boolArr = values.map(b => if b then 1.toByte else 0.toByte) - val buffer = ByteBuffer.allocate(boolArr.length) - buffer.order(ByteOrder.LITTLE_ENDIAN) - buffer.put(boolArr) - val b64String = Base64.getEncoder.encodeToString(buffer.array()) - Tensor(base64Loader(b64String, shape.dimensions.toPythonProxy, "bool")) type Tensor0[V] = Tensor[EmptyTuple, V] type Tensor1[L, V] = Tensor[Tuple1[L], V] @@ -140,55 +108,38 @@ object Tensor0: given int2FloatTensor: Conversion[Int, Tensor0[Float]] = (x: Int) => Tensor0(x.toFloat) given boolean2BooleanTensor: Conversion[Boolean, Tensor0[Boolean]] = (x: Boolean) => Tensor0(x) - def zero[V](vtype: VType[V]): Tensor0[V] = Tensor.zeros(Shape.empty, vtype) - def one[V](vtype: VType[V]): Tensor0[V] = Tensor.ones(Shape.empty, vtype) - def const[V](vtype: VType[V])(value: V)(using writer: Writer[V]): Tensor0[V] = Tensor.const(Shape.empty, vtype)(value) + def apply[V: ExecutionType: Writer](value: V): Tensor0[V] = Tensor(Jax.jnp.full(Shape0.dimensions.toPythonProxy, value, dtype = ExecutionType[V].dtype.jaxType)) + def like[V: Writer](template: Tensor0[V])(value: V): Tensor0[V] = Tensor(Jax.jnp.full(Shape0.dimensions.toPythonProxy, value, dtype = template.dtype.jaxType)) - def randn(key: Random.Key)(using executionType: ExecutionType[Float]): Tensor0[Float] = Normal.standardNormal(Shape.empty).sample(key) def apply[V](jaxValue: Jax.PyDynamic): Tensor0[V] = Tensor(jaxValue) - def apply[V](value: V)(using sv: ExecutionType[V], writer: Writer[V]): Tensor0[V] = Tensor0.const(VType[V])(value) object Tensor1: - def fromArray[L: Label, V](axis: Axis[L], vtype: VType[Float])(values: Array[Float]) = - val dim = (axis -> values.length) - Tensor.fromArray(Shape(dim), vtype)(values) - def fromArray[L: Label, V](axis: Axis[L], vtype: VType[Int])(values: Array[Int]) = - val dim = (axis -> values.length) - Tensor.fromArray(Shape(dim), vtype)(values) - def fromArray[L: Label, V](axis: Axis[L], vtype: VType[Boolean])(values: Array[Boolean]) = - val dim = (axis -> values.length) - Tensor.fromArray(Shape(dim), vtype)(values) + case class Factory[L: Label](val axis: Axis[L]): + private def createShape(l: Int): Shape1[L] = Shape1(axis -> l) + def fromArray[A: ExecutionType, V](values: Array[A])(using t2a: ArrayWriter.Aux[A, V]): Tensor[Tuple1[L], V] = Tensor(createShape(values.length)).fromArray(values) + + def apply[L: Label](axis: Axis[L]): Tensor1.Factory[L] = Tensor1.Factory(axis) object Tensor2: - def fromArray[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], vtype: VType[Float])(values: Array[Array[Float]]): Tensor2[L1, L2, Float] = - val dims = (axis1 -> values.length, axis2 -> values.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten) - def fromArray[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], vtype: VType[Int])(values: Array[Array[Int]]): Tensor2[L1, L2, Int] = - val dims = (axis1 -> values.length, axis2 -> values.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten) - def fromArray[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], vtype: VType[Boolean])(values: Array[Array[Boolean]]): Tensor2[L1, L2, Boolean] = - val dims = (axis1 -> values.length, axis2 -> values.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten) - - def eye[L: Label, V](dim: Dim[L], vtype: VType[V]): Tensor2[L, L, V] = Tensor(Jax.jnp.eye(dim._2, dtype = vtype.dtype.jaxType)) + case class Factory[L1: Label, L2: Label](val axis1: Axis[L1], val axis2: Axis[L2]): + private def createShape[V](valeus: Array[Array[V]]): Shape2[L1, L2] = Shape2(axis1 -> valeus.length, axis2 -> valeus.head.length) + def fromArray[A: ClassTag: ExecutionType, V](values: Array[Array[A]])(using t2a: ArrayWriter.Aux[A, V]): Tensor[(L1, L2), V] = Tensor(createShape(values)).fromArray(values.flatten) + + def apply[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2]): Tensor2.Factory[L1, L2] = Tensor2.Factory(axis1, axis2) + + private def eyeImpl[L: Label, V](dim: Dim[L], dtype: DType): Tensor2[L, L, V] = Tensor(Jax.jnp.eye(dim._2, dtype = dtype.jaxType)) + def eye[L: Label](dim: Dim[L])(using et: ExecutionType[Float]): Tensor2[L, L, Float] = eyeImpl(dim, et.dtype) + def eye[L: Label, V](dim: Dim[L], vtype: VType[V]): Tensor2[L, L, V] = eyeImpl(dim, vtype.dtype) def diag[L: Label, V](diag: Tensor1[L, V]): Tensor2[L, L, V] = Tensor(Jax.jnp.diag(diag.jaxValue)) object Tensor3: - def fromArray[L1: Label, L2: Label, L3: Label, V](axis1: Axis[L1], axis2: Axis[L2], axis3: Axis[L3], vtype: VType[Float])( - values: Array[Array[Array[Float]]] - ): Tensor3[L1, L2, L3, Float] = - val dims = (axis1 -> values.length, axis2 -> values.head.length, axis3 -> values.head.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten.flatten) - def fromArray[L1: Label, L2: Label, L3: Label, V](axis1: Axis[L1], axis2: Axis[L2], axis3: Axis[L3], vtype: VType[Int])( - values: Array[Array[Array[Int]]] - ): Tensor3[L1, L2, L3, Int] = - val dims = (axis1 -> values.length, axis2 -> values.head.length, axis3 -> values.head.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten.flatten) - def fromArray[L1: Label, L2: Label, L3: Label, V](axis1: Axis[L1], axis2: Axis[L2], axis3: Axis[L3], vtype: VType[Boolean])( - values: Array[Array[Array[Boolean]]] - ): Tensor3[L1, L2, L3, Boolean] = - val dims = (axis1 -> values.length, axis2 -> values.head.length, axis3 -> values.head.head.length) - Tensor.fromArray(Shape(dims), vtype)(values.flatten.flatten) + case class Factory[L1: Label, L2: Label, L3: Label](val axis1: Axis[L1], val axis2: Axis[L2], val axis3: Axis[L3]): + private def createShape[V](values: Array[Array[Array[V]]]): Shape3[L1, L2, L3] = + Shape3(axis1 -> values.length, axis2 -> values.head.length, axis3 -> values.head.head.length) + def fromArray[A: ExecutionType: ClassTag, V](values: Array[Array[Array[A]]])(using t2a: ArrayWriter.Aux[A, V]): Tensor3[L1, L2, L3, V] = + Tensor(createShape(values)).fromArray(values.flatten.flatten) + + def apply[L1: Label, L2: Label, L3: Label](axis1: Axis[L1], axis2: Axis[L2], axis3: Axis[L3]): Tensor3.Factory[L1, L2, L3] = Tensor3.Factory(axis1, axis2, axis3) diff --git a/core/src/main/scala/dimwit/tensor/TensorOps.scala b/core/src/main/scala/dimwit/tensor/TensorOps.scala index e29b9515..86d4935e 100644 --- a/core/src/main/scala/dimwit/tensor/TensorOps.scala +++ b/core/src/main/scala/dimwit/tensor/TensorOps.scala @@ -998,13 +998,21 @@ object TensorOps: extension [V: IsNumber: Writer](scalar: V) - def +![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = Tensor0.const(t.vtype)(scalar).broadcastTo(t.shape) + t - def -![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = Tensor0.const(t.vtype)(scalar).broadcastTo(t.shape) - t - def *![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = Tensor0.const(t.vtype)(scalar).broadcastTo(t.shape) * t + def +![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = + given ExecutionType[V] = ExecutionTypeFor[V](t.dtype) + Tensor0(scalar).broadcastTo(t.shape) + t + def -![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = + given ExecutionType[V] = ExecutionTypeFor[V](t.dtype) + Tensor0(scalar).broadcastTo(t.shape) - t + def *![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = + given ExecutionType[V] = ExecutionTypeFor[V](t.dtype) + Tensor0(scalar).broadcastTo(t.shape) * t extension [V: IsFloat: Writer](scalar: V) - def /![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = Tensor0.const(t.vtype)(scalar).broadcastTo(t.shape) / t + def /![T <: Tuple: Labels](t: Tensor[T, V]): Tensor[T, V] = + given ExecutionType[V] = ExecutionTypeFor[V](t.dtype) + Tensor0(scalar).broadcastTo(t.shape) / t object Tensor1Ops: diff --git a/core/src/main/scala/dimwit/tensor/Value.scala b/core/src/main/scala/dimwit/tensor/Value.scala index 535fca56..ab2e5b67 100644 --- a/core/src/main/scala/dimwit/tensor/Value.scala +++ b/core/src/main/scala/dimwit/tensor/Value.scala @@ -2,19 +2,31 @@ package dimwit.tensor import dimwit.stats.Prob import dimwit.stats.LogProb +import scala.compiletime.ops.double +import java.nio.ByteBuffer trait ExecutionType[V]: def dtype: DType object ExecutionType: + def apply[V](using executionType: ExecutionType[V]): ExecutionType[V] = executionType + given floatValue: ExecutionType[Float] with def dtype: DType = DType.Float32 + given intValue: ExecutionType[Int] with def dtype: DType = DType.Int32 + given booleanValue: ExecutionType[Boolean] with def dtype: DType = DType.Bool + given byteValue: ExecutionType[Byte] with + def dtype: DType = DType.Int8 + + given doubleValue: ExecutionType[Double] with + def dtype: DType = DType.Float64 + given prob: ExecutionType[Prob] with def dtype: DType = summon[ExecutionType[Float]].dtype @@ -29,3 +41,5 @@ sealed trait VType[A]: def dtype: DType class OfImpl[A](val dtype: DType) extends VType[A] + +case class ExecutionTypeFor[V](dtype: DType) extends ExecutionType[V] diff --git a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala index aa88ebc7..8e3916f8 100644 --- a/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala +++ b/core/src/test/scala/dimwit/autodiff/AutodiffSuite.scala @@ -27,27 +27,27 @@ class AutodiffSuite extends AnyFunSpec with Matchers: def f(x: Tensor1[A, Float]) = (x * x).sum val df = Autodiff.grad(f) - val x = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 5.0f)) - df(x) shouldEqual Tensor1.fromArray(Axis[A], VType[Float])(Array(2.0f, 10.0f)) + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 5.0f)) + df(x) shouldEqual Tensor1(Axis[A]).fromArray(Array(2.0f, 10.0f)) it("d¹ function using vmap"): def f(x: Tensor2[A, B, Float]) = x.vmap(Axis[A])(_.sum).sum val df = Autodiff.grad(f) - val x = Tensor.ones(Shape(Axis[A] -> 2, Axis[B] -> 2), VType[Float]) - df(x) shouldEqual Tensor.ones(x.shape, x.vtype) + val x = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2)).fill(1f) + df(x) shouldEqual Tensor.like(x).fill(1f) describe("two parameter function"): it("d¹/dx and d¹/dy of (x + 2y)²"): def f(x: Tensor1[A, Float], y: Tensor1[A, Float]) = ((x + (y *! 2.0f)).pow(Tensor0(2.0f))).sum val df = Autodiff.grad(f) - val x = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f)) - val y = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f)) + val x = Tensor1(Axis[A]).fromArray(Array(1.0f)) + val y = Tensor1(Axis[A]).fromArray(Array(1.0f)) val (xGrad, yGrad) = df(x, y).value - xGrad shouldEqual Tensor1.fromArray(Axis[A], VType[Float])(Array(6.0f)) - yGrad shouldEqual Tensor1.fromArray(Axis[A], VType[Float])(Array(12.0f)) + xGrad shouldEqual Tensor1(Axis[A]).fromArray(Array(6.0f)) + yGrad shouldEqual Tensor1(Axis[A]).fromArray(Array(12.0f)) describe("jacobian"): describe("single parameter function"): @@ -55,7 +55,7 @@ class AutodiffSuite extends AnyFunSpec with Matchers: def f(x: Tensor1[A, Float]) = x *! 2.0f val jf = Autodiff.jacobian(f) - val x = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 1.0f)) + val x = Tensor1(Axis[A]).fromArray(Array(1.0f, 1.0f)) jf(x) should approxEqual(Tensor2.eye(x.dim(Axis[A]), x.vtype) *! 2.0f) describe("jacRev / jacFwd"): @@ -71,37 +71,37 @@ class AutodiffSuite extends AnyFunSpec with Matchers: it(s"$modeName d¹ on f: R² -> R², f(x) = swap(x)"): def f(x1: Tensor1[A, Float], x2: Tensor1[A, Float]): (Tensor1[A, Float], Tensor1[A, Float]) = (x2, x1) val df = jacMode(f.tupled) - val x1 = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 0.0f)) - val x2 = Tensor1.fromArray(Axis[A], VType[Float])(Array(0.0f, 1.0f)) + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 0.0f)) + val x2 = Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f)) val (x1Grad, x2Grad) = df(x1, x2) val (x1_dx1, x1_dx2) = x1Grad val (x2_dx1, x2_dx2) = x2Grad - x1_dx1 should approxEqual(Tensor.zeros(x1_dx1.shape, x1_dx1.vtype)) + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) x1_dx2 should approxEqual(Tensor2.eye(x1.dim(Axis[A]), x1.vtype)) x2_dx1 should approxEqual(Tensor2.eye(x2.dim(Axis[A]), x2.vtype)) - x2_dx2 should approxEqual(Tensor.zeros(x2_dx2.shape, x2_dx2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) it(s"$modeName d² on f: R² -> R, f(x1, x2) = sum(x1 * x2)"): def f(x1: Tensor1[A, Float], x2: Tensor1[A, Float]): Tensor0[Float] = (x1 * x2).sum val df = jacMode(f.tupled) val ddf = jacMode(df) - val x1 = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 2.0f)) - val x2 = Tensor1.fromArray(Axis[A], VType[Float])(Array(3.0f, 4.0f)) + val x1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + val x2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) val (x1Grad, x2Grad) = ddf(x1, x2) val (x1_dx1, x1_dx2) = x1Grad val (x2_dx1, x2_dx2) = x2Grad - x1_dx1 should approxEqual(Tensor.zeros(x1_dx1.shape, x1_dx1.vtype)) + x1_dx1 should approxEqual(Tensor.like(x1_dx1).fill(0f)) x1_dx2 should approxEqual(Tensor2.eye(x1.dim(Axis[A]), x1.vtype) *! Tensor0(1.0f)) x2_dx1 should approxEqual(Tensor2.eye(x2.dim(Axis[A]), x2.vtype) *! Tensor0(1.0f)) - x2_dx2 should approxEqual(Tensor.zeros(x2_dx2.shape, x2_dx2.vtype)) + x2_dx2 should approxEqual(Tensor.like(x2_dx2).fill(0f)) describe("Complex application"): it("case class support"): case class Params(w: Tensor1[A, Float], b: Tensor0[Float]) def loss(data: Tensor1[A, Float])(params: Params): Tensor0[Float] = ((data * params.w).sum + params.b).pow(Tensor0(2.0f)) - val trainData = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 2.0f)) + val trainData = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) val dloss = Autodiff.grad(loss(trainData)) - val params = Params(Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 2.0f)), Tensor0(3.0f)) + val params = Params(Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)), Tensor0(3.0f)) val dParams = dloss(params) - dParams.value.w shouldEqual Tensor1.fromArray(Axis[A], VType[Float])(Array(16.0f, 32.0f)) + dParams.value.w shouldEqual Tensor1(Axis[A]).fromArray(Array(16.0f, 32.0f)) diff --git a/core/src/test/scala/dimwit/autodiff/FloatTensorTreeSuite.scala b/core/src/test/scala/dimwit/autodiff/FloatTensorTreeSuite.scala index fb1dfd6e..73d0f5ca 100644 --- a/core/src/test/scala/dimwit/autodiff/FloatTensorTreeSuite.scala +++ b/core/src/test/scala/dimwit/autodiff/FloatTensorTreeSuite.scala @@ -17,9 +17,9 @@ class FloatTensorTreeSuite extends AnyFunSpec with Matchers: val b2: Tensor0[Float] ) val params = Params( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.1f, 0.2f, 0.3f)), + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), Tensor0(0.5f), - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), Tensor0(0.25f) ) val ftTree = summon[FloatTensorTree[Params]] @@ -40,11 +40,11 @@ class FloatTensorTreeSuite extends AnyFunSpec with Matchers: val layer2: LayerParams ) val layer1Params = LayerParams( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), Tensor0(0.25f) ) val layer2Params = LayerParams( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(0.7f, 0.8f), Array(0.9f, 1.0f), Array(1.1f, 1.2f))), + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.7f, 0.8f), Array(0.9f, 1.0f), Array(1.1f, 1.2f))), Tensor0(0.75f) ) val params = ModelParams(layer1Params, layer2Params) @@ -62,7 +62,7 @@ class FloatTensorTreeSuite extends AnyFunSpec with Matchers: val weightBias: (Tensor2[A, B, Float], Tensor0[Float]) ) val layerParams = LayerParams( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f), Array(0.5f, 0.6f))), Tensor0(0.25f) ) val ftTree = summon[FloatTensorTree[LayerParams]] @@ -79,11 +79,11 @@ class FloatTensorTreeSuite extends AnyFunSpec with Matchers: val b1: Tensor0[Float] ) val params1 = Params( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.1f, 0.2f, 0.3f)), + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), Tensor0(0.5f) ) val params2 = Params( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.4f, 0.5f, 0.6f)), + Tensor1(Axis[A]).fromArray(Array(0.4f, 0.5f, 0.6f)), Tensor0(1.5f) ) val ftTree = summon[FloatTensorTree[Params]] diff --git a/core/src/test/scala/dimwit/autodiff/TensorTreeSuite.scala b/core/src/test/scala/dimwit/autodiff/TensorTreeSuite.scala index 575db3a8..30b322a8 100644 --- a/core/src/test/scala/dimwit/autodiff/TensorTreeSuite.scala +++ b/core/src/test/scala/dimwit/autodiff/TensorTreeSuite.scala @@ -16,16 +16,15 @@ class TensorTreeSuite extends AnyFunSpec with Matchers: val flags: Tensor1[A, Boolean] ) val params = Data( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.1f, 0.2f, 0.3f)), - Tensor1.fromArray(Axis[A], VType[Int])(Array(1, 2, 3)), - Tensor1.fromArray(Axis[A], VType[Boolean])(Array(true, false, true)) + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), + Tensor1(Axis[A]).fromArray(Array(1, 2, 3)), + Tensor1(Axis[A]).fromArray(Array(true, false, true)) ) val tree = summon[TensorTree[Data]] - def toOnes[T <: Tuple: Labels, V](t: Tensor[T, V]): Tensor[T, V] = Tensor.ones(t.shape, t.vtype) - val tree2 = tree.map(params, [T <: Tuple, V] => (labels: Labels[T]) ?=> (x: Tensor[T, V]) => toOnes(x)) - tree2.numbers should approxEqual(Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 1.0f, 1.0f))) - tree2.counts should equal(Tensor1.fromArray(Axis[A], VType[Int])(Array(1, 1, 1))) - tree2.flags should equal(Tensor1.fromArray(Axis[A], VType[Boolean])(Array(true, true, true))) + val tree2 = tree.map(params, [T <: Tuple, V] => (labels: Labels[T]) ?=> (x: Tensor[T, V]) => x) + tree2.numbers should approxEqual(params.numbers) + tree2.counts should equal(params.counts) + tree2.flags should equal(params.flags) describe("zipmap"): it("1-level case class"): @@ -34,11 +33,11 @@ class TensorTreeSuite extends AnyFunSpec with Matchers: val b1: Tensor0[Int] ) val params1 = Params( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.1f, 0.2f, 0.3f)), + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), Tensor0(0) ) val params2 = Params( - Tensor1.fromArray(Axis[A], VType[Float])(Array(0.4f, 0.5f, 0.6f)), + Tensor1(Axis[A]).fromArray(Array(0.4f, 0.5f, 0.6f)), Tensor0(1) ) val ftTree = summon[TensorTree[Params]] diff --git a/core/src/test/scala/dimwit/jax/JitSuite.scala b/core/src/test/scala/dimwit/jax/JitSuite.scala index 975366a2..296aa51a 100644 --- a/core/src/test/scala/dimwit/jax/JitSuite.scala +++ b/core/src/test/scala/dimwit/jax/JitSuite.scala @@ -13,7 +13,7 @@ class JitSuite extends AnyFunSpec with Matchers: t * ((t +! 1f) /! 2f) val jitF = jit(f) - val tensor = Tensor.ones(Shape1(Axis[A] -> 5), VType[Float]) + val tensor = Tensor(Shape1(Axis[A] -> 5)).fill(1f) val res = (0 until 25).foldLeft(tensor)((acc, _) => f(acc)) val jittedRes = (0 until 25).foldLeft(tensor)((acc, _) => jitF(acc)) @@ -25,7 +25,7 @@ class JitSuite extends AnyFunSpec with Matchers: t * ((t +! 1f) /! 2f) val (jitDonate, jitF, jitReclaim) = jitDonating(f) - val tensor = Tensor.ones(Shape1(Axis[A] -> 5), VType[Float]) + val tensor = Tensor(Shape1(Axis[A] -> 5)).fill(1f) val res = (0 until 25).foldLeft(tensor)((acc, _) => f(acc)) val jittedRes = jitReclaim((0 until 25).foldLeft(jitDonate(tensor))((acc, _) => jitF(acc))) @@ -37,7 +37,7 @@ class JitSuite extends AnyFunSpec with Matchers: t * ((t +! 1f) /! 2f) val jitF = jitDonatingUnsafe(f) - val tensor = Tensor.ones(Shape1(Axis[A] -> 5), VType[Float]) + val tensor = Tensor(Shape1(Axis[A] -> 5)).fill(1f) val res = (0 until 25).foldLeft(tensor)((acc, _) => f(acc)) val jittedRes = (0 until 25).foldLeft(tensor)((acc, _) => jitF(acc)) @@ -53,7 +53,7 @@ class JitSuite extends AnyFunSpec with Matchers: val end = System.nanoTime() (end - start) / 1_000_000 // ms - val tensor = Tensor.ones(Shape1(Axis[A] -> 5), VType[Float]) + val tensor = Tensor(Shape1(Axis[A] -> 5)).fill(1f) def complexFn(t: Tensor1[A, Float]): Tensor1[A, Float] = (0 until 50).foldLeft(t) { (acc, _) => acc * ((acc +! 1f) /! 2f) } diff --git a/core/src/test/scala/dimwit/random/RandomSuite.scala b/core/src/test/scala/dimwit/random/RandomSuite.scala index df983535..80af9cc8 100644 --- a/core/src/test/scala/dimwit/random/RandomSuite.scala +++ b/core/src/test/scala/dimwit/random/RandomSuite.scala @@ -7,6 +7,7 @@ import me.shadaj.scalapy.py import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers +import dimwit.stats.Normal class RandomSuite extends AnyFunSuite with Matchers: trait A derives Label @@ -42,15 +43,12 @@ class RandomSuite extends AnyFunSuite with Matchers: val key = Random.Key(456) val n = 3 - // Generate random numbers using splitvmap - val vmapResults = key.splitvmap(Axis[Samples] -> n) { k => - Tensor0.randn(k) - } + val vmapResults = key.splitvmap(Axis[Samples] -> n)(Normal.standardSample) // Generate random numbers using individual calls val splitKeys = key.split(n) - val individualResults = Tensor1.fromArray(Axis[Samples], VType[Float])( - splitKeys.map(k => Tensor0.randn(k).item).toArray + val individualResults = Tensor1(Axis[Samples]).fromArray( + splitKeys.map(k => Normal.standardSample(k).item).toArray ) vmapResults should approxEqual(individualResults) @@ -75,7 +73,7 @@ class RandomSuite extends AnyFunSuite with Matchers: // Check that it's actually permuted (with very high probability it won't be identical) // By checking the first element is not 0 (fails 1/10 of the time, but good enough) - val original = Tensor1.fromArray(Axis[A], VType[Int])((0 until n).toArray) + val original = Tensor1(Axis[A]).fromArray((0 until n).toArray) val isIdentity = (perm === original).item isIdentity shouldBe false @@ -85,7 +83,7 @@ class RandomSuite extends AnyFunSuite with Matchers: trait Col derives Label // Create a 2D tensor with distinct values to verify shuffling - val original = Tensor2.fromArray(Axis[Row], Axis[Col], VType[Int])(Array( + val original = Tensor2(Axis[Row], Axis[Col]).fromArray(Array( Array(0, 1, 2), Array(3, 4, 5), Array(6, 7, 8), diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index b70a5a28..c2ba3eef 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/core/src/test/scala/dimwit/stats/DistributionSuite.scala @@ -19,9 +19,9 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("Normal Distribution"): it("logProbs matches JAX"): - val loc = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, 1.0f, -0.5f)) - val scale = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(1.0f, 0.5f, 2.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.5f, -1.0f)) + val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 1.0f, -0.5f)) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, -1.0f)) val dist = Normal(loc, scale) val scalaLogProbs = dist.logProb(x) @@ -32,8 +32,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample means approximates means"): val normal = Normal( - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(0.0f, 1.0f)), - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 0.5f)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 1.0f)), + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => normal.sample(k)) @@ -45,18 +45,18 @@ class DistributionSuite extends AnyFunSpec with Matchers: trait A derives Label trait LocA extends A derives Label trait ScaleA extends A derives Label - val loc = Tensor.fromArray(Shape(Axis[LocA] -> 3), VType[Float])(Array(0.0f, 1.0f, -0.5f)) - val scale = Tensor.fromArray(Shape(Axis[ScaleA] -> 3), VType[Float])(Array(1.0f, 0.5f, 2.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.5f, -1.0f)) + val loc = Tensor(Shape(Axis[LocA] -> 3)).fromArray(Array(0.0f, 1.0f, -0.5f)) + val scale = Tensor(Shape(Axis[ScaleA] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, -1.0f)) val dist = Normal(loc, scale) val scalaLogProbs = dist.logProb(x) scalaLogProbs shouldBe a[Tensor1[A, LogProb]] describe("Uniform Distribution"): it("logProbs matches JAX"): - val low = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, -1.0f, 2.0f)) - val high = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(1.0f, 1.0f, 5.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 0.0f, 3.0f)) + val low = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, -1.0f, 2.0f)) + val high = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 1.0f, 5.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 0.0f, 3.0f)) val dist = Uniform(low, high) val scalaLogProbs = dist.logProb(x) @@ -67,8 +67,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample means approximates means"): val uniform = Uniform( - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(-1.0f, 0.0f)), - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 2.0f)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1.0f, 0.0f)), + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 2.0f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => uniform.sample(k)) @@ -78,8 +78,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("Bernoulli"): it("logProbs matches JAX"): - val probs = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.3f, 0.5f, 0.8f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Int])(Array(0, 1, 1)) + val probs = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.3f, 0.5f, 0.8f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0, 1, 1)) val dist = Bernoulli(probs) val scalaLogProbs = dist.logProb(x) @@ -90,7 +90,7 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample means approximates probabilities"): val bernoulli = Bernoulli( - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(0.3f, 0.7f)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.3f, 0.7f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 1000)(k => bernoulli.sample(k)) @@ -100,9 +100,9 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("Cauchy"): it("logProbs matches JAX"): - val loc = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, 1.0f, -0.5f)) - val scale = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(1.0f, 0.5f, 2.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.5f, -1.0f)) + val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 1.0f, -0.5f)) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, -1.0f)) val dist = Cauchy(loc, scale) val scalaLogProbs = dist.logProb(x) @@ -113,8 +113,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample medians approximates location"): val cauchy = Cauchy( - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(0.0f, 2.0f)), - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 0.5f)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 2.0f)), + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 50000)(k => cauchy.sample(k)) @@ -124,9 +124,9 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("HalfNormal"): it("logProbs computed correctly"): - val loc = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, 0.0f, 0.0f)) - val scale = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(1.0f, 0.5f, 2.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.0f, 0.8f)) + val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 0.0f, 0.0f)) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.0f, 0.8f)) val dist = HalfNormal(loc, scale) val scalaLogProbs = dist.logProb(x) @@ -138,8 +138,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample means approximates expected means"): val halfNormal = HalfNormal( - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(0.0f, 0.0f)), - Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 2.0f)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 0.0f)), + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 2.0f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => halfNormal.sample(k)) @@ -152,9 +152,9 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("StudentT"): it("logProbs matches JAX"): val df = 5 - val loc = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, 1.0f, -0.5f)) - val scale = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(1.0f, 0.5f, 2.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.5f, -1.0f)) + val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 1.0f, -0.5f)) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, -1.0f)) val dist = StudentT(df, loc, scale) val scalaLogProbs = dist.logProb(x) @@ -166,8 +166,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: it("sample means approximates location"): val studentT = StudentT( df = 5, - loc = Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(0.0f, 2.0f)), - scale = Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 0.5f)) + loc = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 2.0f)), + scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => studentT.sample(k)) @@ -177,15 +177,15 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("MVNormal"): it("logProb matches JAX"): - val mean = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.0f, 1.0f, 2.0f)) - val cov = Tensor.fromArray(Shape(Axis[A] -> 3, Axis[Prime[A]] -> 3), VType[Float])( + val mean = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 1.0f, 2.0f)) + val cov = Tensor(Shape(Axis[A] -> 3, Axis[Prime[A]] -> 3)).fromArray( Array( 1.0f, 0.5f, 0.2f, 0.5f, 2.0f, 0.3f, 0.2f, 0.3f, 1.5f ) ) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.5f, 1.5f, 2.2f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 1.5f, 2.2f)) val dist = MVNormal(mean, cov) val scalaLogProb = dist.logProb(x) @@ -195,8 +195,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: scalaLogProb.asFloat should approxEqual(jaxLogProb) it("sample mean approximates mean"): - val mean = Tensor.fromArray(Shape(Axis[A] -> 2), VType[Float])(Array(1.0f, 2.0f)) - val cov = Tensor.fromArray(Shape(Axis[A] -> 2, Axis[Prime[A]] -> 2), VType[Float])( + val mean = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 2.0f)) + val cov = Tensor(Shape(Axis[A] -> 2, Axis[Prime[A]] -> 2)).fromArray( Array(1.0f, 0.3f, 0.3f, 1.0f) ) val mvNormal = MVNormal(mean, cov) @@ -208,8 +208,8 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("Dirichlet"): it("logProb matches JAX"): - val concentration = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(2.0f, 3.0f, 5.0f)) - val x = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.2f, 0.3f, 0.5f)) + val concentration = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(2.0f, 3.0f, 5.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.2f, 0.3f, 0.5f)) val dist = Dirichlet(concentration) val scalaLogProb = dist.logProb(x) @@ -219,21 +219,21 @@ class DistributionSuite extends AnyFunSpec with Matchers: scalaLogProb.asFloat should approxEqual(jaxLogProb) it("sample mean approximates expected mean"): - val concentration = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(2.0f, 5.0f, 3.0f)) + val concentration = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(2.0f, 5.0f, 3.0f)) val dirichlet = Dirichlet(concentration) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => dirichlet.sample(k)) val sampleMean = samples.mean(Axis[Samples]) // Expected mean for Dirichlet is concentration / sum(concentration) // For [2.0, 5.0, 3.0], sum=10.0, so expected is [0.2, 0.5, 0.3] - val expectedMean = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.2f, 0.5f, 0.3f)) + val expectedMean = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.2f, 0.5f, 0.3f)) sampleMean should approxEqual(expectedMean, 0.2f) describe("Multinomial"): it("logProb matches JAX"): - val probsFloat = Tensor.fromArray(Shape(Axis[A] -> 4), VType[Float])(Array(0.1f, 0.2f, 0.3f, 0.4f)) + val probsFloat = Tensor(Shape(Axis[A] -> 4)).fromArray(Array(0.1f, 0.2f, 0.3f, 0.4f)) val probs = Prob(probsFloat) - val x = Tensor.fromArray(Shape(Axis[A] -> 4), VType[Int])(Array(2, 1, 3, 4)) + val x = Tensor(Shape(Axis[A] -> 4)).fromArray(Array(2, 1, 3, 4)) val n = 10 val dist = Multinomial[A](n, probs) @@ -244,7 +244,7 @@ class DistributionSuite extends AnyFunSpec with Matchers: scalaLogProb.asFloat should approxEqual(jaxLogProb) it("sample mean approximates expected counts"): - val probsFloat = Tensor.fromArray(Shape(Axis[A] -> 3), VType[Float])(Array(0.2f, 0.5f, 0.3f)) + val probsFloat = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.2f, 0.5f, 0.3f)) val probs = Prob(probsFloat) val n = 100 val multinomial = Multinomial[A](n, probs) @@ -257,7 +257,7 @@ class DistributionSuite extends AnyFunSpec with Matchers: describe("Categorical"): it("logProb matches expected value"): - val probs = Tensor.fromArray(Shape(Axis[A] -> 4), VType[Float])(Array(0.1f, 0.2f, 0.3f, 0.4f)) + val probs = Tensor(Shape(Axis[A] -> 4)).fromArray(Array(0.1f, 0.2f, 0.3f, 0.4f)) val x = Tensor0(2) val dist = Categorical(probs) @@ -266,7 +266,7 @@ class DistributionSuite extends AnyFunSpec with Matchers: scalaLogProb.asFloat should approxEqual(expectedLogProb) it("sample distribution matches probabilities"): - val probs = Tensor.fromArray(Shape(Axis[A] -> 4), VType[Float])(Array(0.1f, 0.2f, 0.3f, 0.4f)) + val probs = Tensor(Shape(Axis[A] -> 4)).fromArray(Array(0.1f, 0.2f, 0.3f, 0.4f)) val categorical = Categorical(probs) val key = Random.Key(42) val numSamples = 10000 diff --git a/core/src/test/scala/dimwit/tensor/TensorCompileSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCompileSuite.scala index e4ca2d48..05eef31e 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCompileSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCompileSuite.scala @@ -9,7 +9,7 @@ import scala.compiletime.testing.typeCheckErrors class TensorCompileSuite extends AnyFunSpec with Matchers: it("Nice error message when axis not found in tensor for sum"): - val t = Tensor.zeros(Shape(Axis[A] -> 1, Axis[B] -> 2), VType[Float]) + val t = Tensor(Shape(Axis[A] -> 1, Axis[B] -> 2)).fill(0f) // val res = t.sum(Axis[C]) val errors = typeCheckErrors("t.sum(Axis[C])") errors should have size 1 @@ -17,7 +17,7 @@ class TensorCompileSuite extends AnyFunSpec with Matchers: error.message should include("Axis[dimwit.C] not found in Tensor[(dimwit.A, dimwit.B)]") it("Nice error message when axes not found in tensor for sum"): - val t = Tensor.zeros(Shape(Axis[A] -> 1, Axis[B] -> 2), VType[Float]) + val t = Tensor(Shape(Axis[A] -> 1, Axis[B] -> 2)).fill(0f) // val res = t.sum((Axis[A], Axis[C])) val errors = typeCheckErrors("t.sum((Axis[A], Axis[C]))") errors should have size 1 @@ -25,8 +25,8 @@ class TensorCompileSuite extends AnyFunSpec with Matchers: error.message should include("(dimwit.tensor.Axis[dimwit.A], dimwit.tensor.Axis[dimwit.C])]] not all found in Tensor shape [(dimwit.A, dimwit.B)]") it("Nice error message when axes not found in zipvmap"): - val ab = Tensor.zeros(Shape(Axis[A] -> 1, Axis[B] -> 2), VType[Float]) - val bc = Tensor.zeros(Shape(Axis[B] -> 2, Axis[C] -> 1), VType[Float]) + val ab = Tensor(Shape(Axis[A] -> 1, Axis[B] -> 2)).fill(0f) + val bc = Tensor(Shape(Axis[B] -> 2, Axis[C] -> 1)).fill(0f) // val res = zipvmap(Axis[C])(ab, bc) { case (x, y) => x.sum + y.sum } val errors = typeCheckErrors("zipvmap(Axis[C])(ab, bc) { case (x, y) => x.sum + y.sum }") errors should have size 1 diff --git a/core/src/test/scala/dimwit/tensor/TensorCovarianceSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCovarianceSuite.scala index b2912bc0..eb6a371d 100644 --- a/core/src/test/scala/dimwit/tensor/TensorCovarianceSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorCovarianceSuite.scala @@ -14,9 +14,9 @@ class TensorCovarianceSuite extends AnyFunSpec with Matchers: trait Child2 extends Parent derives Label trait NoChild derives Label def concreteFunction(t: Tensor1[Parent, Float]): Tensor1[Parent, Float] = t + t - val child1: Tensor1[Child1, Float] = Tensor.ones(Shape1(Axis[Child1] -> 4), VType[Float]) - val child2: Tensor1[Child2, Float] = Tensor.ones(Shape1(Axis[Child2] -> 4), VType[Float]) - val noChild: Tensor1[NoChild, Float] = Tensor.ones(Shape1(Axis[NoChild] -> 4), VType[Float]) + val child1: Tensor1[Child1, Float] = Tensor(Shape1(Axis[Child1] -> 4)).fill(1f) + val child2: Tensor1[Child2, Float] = Tensor(Shape1(Axis[Child2] -> 4)).fill(1f) + val noChild: Tensor1[NoChild, Float] = Tensor(Shape1(Axis[NoChild] -> 4)).fill(1f) "concreteFunction(child1)" should compile "concreteFunction(child2)" should compile @@ -27,8 +27,8 @@ class TensorCovarianceSuite extends AnyFunSpec with Matchers: trait Child1 extends Parent derives Label trait Child2 extends Parent derives Label def genericFunction[T <: Parent: Label](t: Tensor1[T, Float]): Tensor1[T, Float] = t + t - val child1: Tensor1[Child1, Float] = Tensor.ones(Shape1(Axis[Child1] -> 4), VType[Float]) - val child2: Tensor1[Child2, Float] = Tensor.ones(Shape1(Axis[Child2] -> 4), VType[Float]) + val child1: Tensor1[Child1, Float] = Tensor(Shape1(Axis[Child1] -> 4)).fill(1f) + val child2: Tensor1[Child2, Float] = Tensor(Shape1(Axis[Child2] -> 4)).fill(1f) "genericFunction(child1)" should compile "genericFunction(child2)" should compile @@ -41,8 +41,8 @@ class TensorCovarianceSuite extends AnyFunSpec with Matchers: opaque type Logit = Float opaque type Prob = Float - def createLogits[L: Label](s: Shape1[L]): Tensor1[L, Logit] = Tensor.zeros(s, VType[Logit]) - def createProbs[L: Label](s: Shape1[L]): Tensor1[L, Prob] = Tensor.zeros(s, VType[Prob]) + def createLogits[L: Label](s: Shape1[L]): Tensor1[L, Logit] = Tensor(s).fill(0f) + def createProbs[L: Label](s: Shape1[L]): Tensor1[L, Prob] = Tensor(s).fill(0f) // Operation restricted only to Logit 'land' def combineLogits[L: Label](a: Tensor1[L, Logit], b: Tensor1[L, Logit]): Tensor1[L, Logit] = a + b @@ -52,7 +52,7 @@ class TensorCovarianceSuite extends AnyFunSpec with Matchers: val shape = Shape1(Axis[Classes] -> 10) val logits = MLContext.createLogits(shape) val probs = MLContext.createProbs(shape) - val rawFloats = Tensor.ones(shape, VType[Float]) + val rawFloats = Tensor(shape).fill(1f) "MLContext.combineLogits(logits, logits)" should compile "MLContext.combineProbs(probs, probs)" should compile diff --git a/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala new file mode 100644 index 00000000..a2211d65 --- /dev/null +++ b/core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala @@ -0,0 +1,62 @@ +package dimwit.tensor + +import dimwit.* +import org.scalatest.propspec.AnyPropSpec +import org.scalatest.matchers.should.Matchers +import org.scalatest.funspec.AnyFunSpec +import scala.compiletime.testing.typeCheckErrors + +class TensorCreationSuite extends AnyFunSpec with Matchers: + + def withJaxX64Support[R](block: => R): R = + import me.shadaj.scalapy.py + val jaxConfig = py.module("jax").config + val current = jaxConfig.jax_enable_x64.as[Boolean] + jaxConfig.update("jax_enable_x64", true) + val res = block + jaxConfig.update("jax_enable_x64", current) + res + + describe("Default settings"): + describe("Tensor fill"): + it("Fill tensors with tensor types"): + val intTensor = Tensor(Shape2(Axis[A] -> 4, Axis[B] -> 5)).fill(42) + intTensor.dtype shouldBe DType.Int32 + val floatTensor = Tensor(Shape3(Axis[A] -> 2, Axis[B] -> 3, Axis[C] -> 4)).fill(3.14f) + floatTensor.dtype shouldBe DType.Float32 + val boolTensor = Tensor(Shape1(Axis[A] -> 10)).fill(true) + boolTensor.dtype shouldBe DType.Bool + + it("Fill tensors with widened types"): + // Test byte defaults to int8 + val intTensorFromByte = Tensor(Shape2(Axis[A] -> 4, Axis[B] -> 5)).fill(42.toByte) + intTensorFromByte.dtype shouldBe DType.Int8 + // Test double defaults to float64 + withJaxX64Support: // Enable float64 support in JAX + val floatTensorFromDouble = Tensor(Shape3(Axis[A] -> 2, Axis[B] -> 3, Axis[C] -> 4)).fill(3.14) + floatTensorFromDouble.dtype shouldBe DType.Float64 + describe("Tensor fromArray"): + it("fromArray with tensor types"): + val intTensor = Tensor(Shape1(Axis[A] -> 3)).fromArray(Array(1, 2, 3)) + intTensor.dtype shouldBe DType.Int32 + val floatTensor = Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 2)).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) + floatTensor.dtype shouldBe DType.Float32 + it("fromArray with widened types"): + // Test short defaults to int8 + val intTensorFromShort = Tensor(Shape1(Axis[A] -> 3)).fromArray(Array(1.toByte, 2.toByte, 3.toByte)) + intTensorFromShort.dtype shouldBe DType.Int8 + // Test double defaults to float64 + withJaxX64Support: // Enable float64 support in JAX + val floatTensorFromDouble = Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 2)).fromArray(Array(1.0, 2.0, 3.0, 4.0)) + floatTensorFromDouble.dtype shouldBe DType.Float64 + + describe("Overwrite default setings"): + it("Change double default dtype from Float64 to Float32"): + given ExecutionType[Double] = ExecutionTypeFor[Double](DType.Float32) + // Check fill + val floatTensorFromDouble = Tensor(Shape3(Axis[A] -> 2, Axis[B] -> 3, Axis[C] -> 4)).fill(3.14) + floatTensorFromDouble.dtype shouldBe DType.Float32 + // Check fromArray + withJaxX64Support: // Enable float64 support in JAX + val floatTensorFromDouble2 = Tensor(Shape2(Axis[A] -> 2, Axis[B] -> 2)).fromArray(Array(1.0, 2.0, 3.0, 4.0)) + floatTensorFromDouble2.dtype shouldBe DType.Float32 diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsBinarySuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsBinarySuite.scala index e2f720a4..d4699e69 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsBinarySuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsBinarySuite.scala @@ -7,51 +7,51 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsBinarySuite extends AnyFunSpec with Matchers: - val t2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(10.0f, 20.0f), Array(30.0f, 40.0f)) ) - val t2_2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2_2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(2.0f, 4.0f), Array(5.0f, 8.0f)) ) - val i2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Int])( + val i2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(10, 20), Array(30, 40)) ) - val i2_2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Int])( + val i2_2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(2, 4), Array(5, 8)) ) describe("Float Binary Ops"): it("Addition (+)"): - (t2 + t2_2) shouldEqual Tensor.fromArray(t2.shape, t2.vtype)(Array(12.0f, 24.0f, 35.0f, 48.0f)) + (t2 + t2_2) shouldEqual Tensor.like(t2).fromArray(Array(12.0f, 24.0f, 35.0f, 48.0f)) it("Subtraction (-)"): - (t2 - t2_2) shouldEqual Tensor.fromArray(t2.shape, t2.vtype)(Array(8.0f, 16.0f, 25.0f, 32.0f)) + (t2 - t2_2) shouldEqual Tensor.like(t2).fromArray(Array(8.0f, 16.0f, 25.0f, 32.0f)) it("Multiplication (*)"): - (t2 * t2_2) shouldEqual Tensor.fromArray(t2.shape, t2.vtype)(Array(20.0f, 80.0f, 150.0f, 320.0f)) + (t2 * t2_2) shouldEqual Tensor.like(t2).fromArray(Array(20.0f, 80.0f, 150.0f, 320.0f)) it("Division (/)"): - (t2 / t2_2) shouldEqual Tensor.fromArray(t2.shape, t2.vtype)(Array(5.0f, 5.0f, 6.0f, 5.0f)) + (t2 / t2_2) shouldEqual Tensor.like(t2).fromArray(Array(5.0f, 5.0f, 6.0f, 5.0f)) it("Comparisons (<, <=, >, >=)"): - (t2 < t2_2).asBoolean shouldEqual Tensor.fromArray(t2.shape, VType[Boolean])(Array(false, false, false, false)) - (t2 > t2_2).asBoolean shouldEqual Tensor.fromArray(t2.shape, VType[Boolean])(Array(true, true, true, true)) + (t2 < t2_2).asBoolean shouldEqual Tensor(t2.shape).fromArray(Array(false, false, false, false)) + (t2 > t2_2).asBoolean shouldEqual Tensor(t2.shape).fromArray(Array(true, true, true, true)) it("elementEquals"): - (t2 `elementEquals` t2) shouldEqual Tensor.fromArray(t2.shape, VType[Boolean])(Array(true, true, true, true)) - (t2 `elementEquals` t2_2) shouldEqual Tensor.fromArray(t2.shape, VType[Boolean])(Array(false, false, false, false)) + (t2 `elementEquals` t2) shouldEqual Tensor(t2.shape).fromArray(Array(true, true, true, true)) + (t2 `elementEquals` t2_2) shouldEqual Tensor(t2.shape).fromArray(Array(false, false, false, false)) describe("Int Binary Ops"): it("Addition (+)"): - (i2 + i2_2) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(12, 24, 35, 48)) + (i2 + i2_2) shouldEqual Tensor.like(i2).fromArray(Array(12, 24, 35, 48)) it("Subtraction (-)"): - (i2 - i2_2) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(8, 16, 25, 32)) + (i2 - i2_2) shouldEqual Tensor.like(i2).fromArray(Array(8, 16, 25, 32)) it("Multiplication (*)"): - (i2 * i2_2) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(20, 80, 150, 320)) + (i2 * i2_2) shouldEqual Tensor.like(i2).fromArray(Array(20, 80, 150, 320)) it("Comparisons (<, <=, >, >=)"): - (i2 < i2_2).asBoolean shouldEqual Tensor.fromArray(i2.shape, VType[Boolean])(Array(false, false, false, false)) - (i2 >= i2_2).asBoolean shouldEqual Tensor.fromArray(i2.shape, VType[Boolean])(Array(true, true, true, true)) + (i2 < i2_2).asBoolean shouldEqual Tensor(i2.shape).fromArray(Array(false, false, false, false)) + (i2 >= i2_2).asBoolean shouldEqual Tensor(i2.shape).fromArray(Array(true, true, true, true)) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala index a44493e6..b45b0d57 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala @@ -7,28 +7,27 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: - val tA = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 2.0f)) + val tA = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) - val tAB = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(10.0f, 20.0f), Array(30.0f, 40.0f))) + val tAB = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(10.0f, 20.0f), Array(30.0f, 40.0f))) - val tAB2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(100.0f, 200.0f))) - - val iA = Tensor1.fromArray(Axis[A], VType[Int])(Array(1, 2)) - val iAB = Tensor2.fromArray(Axis[A], Axis[B], VType[Int])(Array(Array(1, 2), Array(3, 4))) + val tAB2 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(100.0f, 200.0f))) + val iA = Tensor1(Axis[A]).fromArray(Array(1, 2)) + val iAB = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1, 2), Array(3, 4))) describe("Scalar Broadcasting"): describe("Int"): it("Addition"): - (5 +! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(6, 7, 8, 9)) + (5 +! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(6, 7, 8, 9)) (5 +! iAB) shouldEqual (iAB +! 5) it("Subtraction"): - (5 -! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(4, 3, 2, 1)) - (iAB -! 5) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(-4, -3, -2, -1)) + (5 -! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(4, 3, 2, 1)) + (iAB -! 5) shouldEqual Tensor.like(iAB).fromArray(Array(-4, -3, -2, -1)) it("Multiplication"): - (3 *! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(3, 6, 9, 12)) + (3 *! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(3, 6, 9, 12)) (3 *! iAB) shouldEqual (iAB *! 3) it("No Int Division Supported"): @@ -38,35 +37,35 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: describe("Float"): it("Addition"): - (2.0f +! tAB) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(12.0f, 22.0f, 32.0f, 42.0f)) + (2.0f +! tAB) shouldEqual Tensor.like(tAB).fromArray(Array(12.0f, 22.0f, 32.0f, 42.0f)) (2.0f +! tAB) shouldEqual (tAB +! 2.0f) it("Subtraction"): - (5.0f -! tAB) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(-5.0f, -15.0f, -25.0f, -35.0f)) - (tAB -! 5.0f) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(5.0f, 15.0f, 25.0f, 35.0f)) + (5.0f -! tAB) shouldEqual Tensor.like(tAB).fromArray(Array(-5.0f, -15.0f, -25.0f, -35.0f)) + (tAB -! 5.0f) shouldEqual Tensor.like(tAB).fromArray(Array(5.0f, 15.0f, 25.0f, 35.0f)) it("Multiplication"): - (2.0f *! tAB) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(20.0f, 40.0f, 60.0f, 80.0f)) + (2.0f *! tAB) shouldEqual Tensor.like(tAB).fromArray(Array(20.0f, 40.0f, 60.0f, 80.0f)) (2.0f *! tAB) shouldEqual (tAB *! 2.0f) it("Division"): - (2.0f /! tAB) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(0.2f, 0.1f, 0.06666667f, 0.05f)) - (tAB /! 2.0f) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)(Array(5.0f, 10.0f, 15.0f, 20.0f)) + (2.0f /! tAB) shouldEqual Tensor.like(tAB).fromArray(Array(0.2f, 0.1f, 0.06666667f, 0.05f)) + (tAB /! 2.0f) shouldEqual Tensor.like(tAB).fromArray(Array(5.0f, 10.0f, 15.0f, 20.0f)) describe("Vector-to-Tensor Broadcasting"): describe("Int"): it("Addition"): - (iA +! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(2, 3, 5, 6)) + (iA +! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(2, 3, 5, 6)) (iA +! iAB) shouldEqual (iAB +! iA) it("Subtraction"): - (iA -! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(0, -1, -1, -2)) - (iAB -! iA) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(0, 1, 1, 2)) + (iA -! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(0, -1, -1, -2)) + (iAB -! iA) shouldEqual Tensor.like(iAB).fromArray(Array(0, 1, 1, 2)) it("Multiplication"): - (iA *! iAB) shouldEqual Tensor.fromArray(iAB.shape, iAB.vtype)(Array(1, 2, 6, 8)) + (iA *! iAB) shouldEqual Tensor.like(iAB).fromArray(Array(1, 2, 6, 8)) (iA *! iAB) shouldEqual (iAB *! iA) it("No Int Division Supported"): @@ -77,7 +76,7 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: it("Addition"): (tAB +! tA) should approxEqual( - Tensor.fromArray(tAB.shape, tAB.vtype)( + Tensor.like(tAB).fromArray( Array(11.0f, 21.0f, 32.0f, 42.0f) ) ) @@ -85,17 +84,17 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: it("Subtraction"): (tAB -! tA) should approxEqual( - Tensor.fromArray(tAB.shape, tAB.vtype)( + Tensor.like(tAB).fromArray( Array(9.0f, 19.0f, 28.0f, 38.0f) ) ) - (tA -! tAB) shouldEqual Tensor.fromArray(tAB.shape, tAB.vtype)( + (tA -! tAB) shouldEqual Tensor.like(tAB).fromArray( Array(-9.0f, -19.0f, -28.0f, -38.0f) ) it("Multiplication"): (tAB *! tA) should approxEqual( - Tensor.fromArray(tAB.shape, tAB.vtype)( + Tensor.like(tAB).fromArray( Array(10.0f, 20.0f, 60.0f, 80.0f) ) ) @@ -103,24 +102,24 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: it("Division"): (tAB /! tA) should approxEqual( - Tensor.fromArray(tAB.shape, tAB.vtype)( + Tensor.like(tAB).fromArray( Array(10.0f, 20.0f, 15.0f, 20.0f) ) ) (tA /! tAB) should approxEqual( - Tensor.fromArray(tAB.shape, tAB.vtype)( + Tensor.like(tAB).fromArray( Array(0.1f, 0.05f, 0.06666667f, 0.05f) ) ) describe("Tensor-to-Tensor Broadcasting (complex)"): - val tABCD = Tensor.fromArray(Shape(Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2), VType[Float])( + val tABCD = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2)).fromArray( Array.range(1, 17).map(_.toFloat) ) it("AB broadcastTo ABCD"): - val AB = Tensor.fromArray(Shape(Axis[A] -> 2, Axis[B] -> 2), VType[Float])( + val AB = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2)).fromArray( Array.range(1, 5).map(_.toFloat) ) val res = AB.broadcastTo(tABCD.shape) @@ -131,7 +130,7 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: res.slice((Axis[C] -> 1, Axis[D] -> 1)) should approxEqual(AB) it("BC broadcastTo ABCD"): - val BC = Tensor.fromArray(Shape(Axis[B] -> 2, Axis[C] -> 2), VType[Float])( + val BC = Tensor(Shape(Axis[B] -> 2, Axis[C] -> 2)).fromArray( Array.range(1, 5).map(_.toFloat) ) val res = BC.broadcastTo(tABCD.shape) @@ -142,7 +141,7 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: res.slice((Axis[A] -> 1, Axis[D] -> 1)) should approxEqual(BC) it("CD broadcastTo ABCD"): - val CD = Tensor.fromArray(Shape(Axis[C] -> 2, Axis[D] -> 2), VType[Float])( + val CD = Tensor(Shape(Axis[C] -> 2, Axis[D] -> 2)).fromArray( Array.range(1, 5).map(_.toFloat) ) val res = CD.broadcastTo(tABCD.shape) @@ -154,7 +153,7 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: describe("Disallow"): - val tABCD = Tensor.fromArray(Shape(Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2), VType[Float])( + val tABCD = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2)).fromArray( Array.range(1, 17).map(_.toFloat) ) @@ -171,7 +170,7 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: describe("Operator Precedence"): it("multiplication (*!) binds tighter than addition (+!)"): - val tA = Tensor.ones(Shape1(tAB.shape.dim(Axis[A])), VType[Float]) + val tA = Tensor(Shape1(tAB.shape.dim(Axis[A]))).fill(1f) val res = tAB *! Tensor0(2.0f) +! tA val correct = (tAB *! Tensor0(2.0f)) +! tA val wrong = tAB *! (Tensor0(2.0f) +! tA) @@ -181,6 +180,6 @@ class TensorOpsBroadcastSuite extends AnyFunSpec with Matchers: describe("Mixed Broadcasting Cases"): it("Broadcasting ab + bc to abc"): - val ab = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(1.0f, 2.0f))) - val bc = Tensor2.fromArray(Axis[B], Axis[C], VType[Float])(Array(Array(10.0f), Array(20.0f))) + val ab = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) + val bc = Tensor2(Axis[B], Axis[C]).fromArray(Array(Array(10.0f), Array(20.0f))) "ab +! bc" shouldNot compile // TODO add support for this diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala index a9ab560e..fb5113ce 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsContractionSuite.scala @@ -7,18 +7,18 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsContractionSuite extends AnyFunSpec with Matchers: - val v1 = Tensor1.fromArray(Axis[A], VType[Float])( + val v1 = Tensor1(Axis[A]).fromArray( Array(1.0f, 2.0f) ) - val v2 = Tensor1.fromArray(Axis[A], VType[Float])( + val v2 = Tensor1(Axis[A]).fromArray( Array(3.0f, 4.0f) ) - val m1 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val m1 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f)) ) - val m2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val m2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(10.0f, 20.0f), Array(30.0f, 40.0f)) ) @@ -32,7 +32,7 @@ class TensorOpsContractionSuite extends AnyFunSpec with Matchers: res.shape.labels shouldBe List("B", "B'") res should approxEqual( - Tensor.fromArray(res.shape, res.vtype)( + Tensor.like(res).fromArray( Array(100.0f, 140.0f, 140.0f, 200.0f) ) ) @@ -42,7 +42,7 @@ class TensorOpsContractionSuite extends AnyFunSpec with Matchers: res.shape.labels shouldBe List("A", "A'") res should approxEqual( - Tensor.fromArray(res.shape, res.vtype)( + Tensor.like(res).fromArray( Array(50.0f, 110.0f, 110.0f, 250.0f) ) ) @@ -55,7 +55,7 @@ class TensorOpsContractionSuite extends AnyFunSpec with Matchers: res.shape.labels shouldBe List("B", "D") res should approxEqual( - Tensor.fromArray(res.shape, res.vtype)( + Tensor.like(res).fromArray( Array(100.0f, 140.0f, 140.0f, 200.0f) ) ) @@ -67,12 +67,12 @@ class TensorOpsContractionSuite extends AnyFunSpec with Matchers: describe("outerProduct"): it("Tensor1[A] and Tensor1[B] to Tensor2[A, B]"): - val vA = Tensor1.fromArray(Axis[A], VType[Float])(Array(1.0f, 2.0f)) - val vB = Tensor1.fromArray(Axis[B], VType[Float])(Array(10.0f, 20.0f)) + val vA = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) + val vB = Tensor1(Axis[B]).fromArray(Array(10.0f, 20.0f)) val res = vA.outerProduct(vB) res should approxEqual( - Tensor.fromArray(res.shape, res.vtype)( + Tensor.like(res).fromArray( Array(10.0f, 20.0f, 20.0f, 40.0f) ) ) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala index 95d3acd0..97208d6d 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala @@ -7,21 +7,21 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsElementwiseSuite extends AnyFunSpec with Matchers: - val t2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2 = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(-1.0f, 0.0f), Array(1.0f, 4.0f) ) ) - val i2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Int])( + val i2 = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(-1, 0), Array(1, 2) ) ) - val b2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Boolean])( + val b2 = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(true, false), Array(false, true) @@ -31,34 +31,35 @@ class TensorOpsElementwiseSuite extends AnyFunSpec with Matchers: describe("Float ops (Tensor2)"): it("abs"): - t2.abs should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(1.0f, 0.0f, 1.0f, 4.0f))) + t2.abs should approxEqual(Tensor.like(t2).fromArray(Array(1.0f, 0.0f, 1.0f, 4.0f))) it("sign"): - t2.sign should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(-1.0f, 0.0f, 1.0f, 1.0f))) + t2.sign should approxEqual(Tensor.like(t2).fromArray(Array(-1.0f, 0.0f, 1.0f, 1.0f))) it("pow"): - t2.pow(Tensor0(2.0f)) should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(1.0f, 0.0f, 1.0f, 16.0f))) + t2.pow(Tensor0(2.0f)) should approxEqual(Tensor.like(t2).fromArray(Array(1.0f, 0.0f, 1.0f, 16.0f))) it("sqrt"): - val tPos = Tensor.fromArray(t2.shape, t2.vtype)(Array(4.0f, 9.0f, 16.0f, 25.0f)) - tPos.sqrt should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(2.0f, 3.0f, 4.0f, 5.0f))) + val tPos = Tensor.like(t2).fromArray(Array(4.0f, 9.0f, 16.0f, 25.0f)) + tPos.sqrt should approxEqual(Tensor.like(t2).fromArray(Array(2.0f, 3.0f, 4.0f, 5.0f))) it("exp/log (identity)"): - val tZero = Tensor.zeros(t2.shape, t2.vtype) - tZero.exp should approxEqual(Tensor.ones(t2.shape, t2.vtype)) - Tensor.ones(t2.shape, t2.vtype).log should approxEqual(tZero) + val tZero = Tensor.like(t2).fill(0f) + val tOne = Tensor.like(t2).fill(1f) + tZero.exp should approxEqual(tOne) + tOne.log should approxEqual(tZero) it("sin/cos/tanh"): - val tZero = Tensor.zeros(t2.shape, t2.vtype) + val tZero = Tensor.like(t2).fill(0f) tZero.sin should approxEqual(tZero) - tZero.cos should approxEqual(Tensor.ones(t2.shape, t2.vtype)) + tZero.cos should approxEqual(Tensor.like(t2).fill(1f)) tZero.tanh should approxEqual(tZero) it("clip"): - t2.clip(0.0f, 2.0f) should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(0.0f, 0.0f, 1.0f, 2.0f))) + t2.clip(0.0f, 2.0f) should approxEqual(Tensor.like(t2).fromArray(Array(0.0f, 0.0f, 1.0f, 2.0f))) it("unary_-"): - (-t2) should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)(Array(1.0f, 0.0f, -1.0f, -4.0f))) + (-t2) should approxEqual(Tensor.like(t2).fromArray(Array(1.0f, 0.0f, -1.0f, -4.0f))) it("approxEquals / approxElementEquals"): val t2Near = t2 *! Tensor0(1.0000001f) @@ -68,24 +69,24 @@ class TensorOpsElementwiseSuite extends AnyFunSpec with Matchers: describe("Int ops (Tensor2)"): it("abs"): - i2.abs shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(1, 0, 1, 2)) + i2.abs shouldEqual Tensor.like(i2).fromArray(Array(1, 0, 1, 2)) it("sign"): - i2.sign shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(-1, 0, 1, 1)) + i2.sign shouldEqual Tensor.like(i2).fromArray(Array(-1, 0, 1, 1)) it("pow"): - i2.pow(Tensor0(3)) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(-1, 0, 1, 8)) + i2.pow(Tensor0(3)) shouldEqual Tensor.like(i2).fromArray(Array(-1, 0, 1, 8)) it("clip"): - i2.clip(0, 1) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(0, 0, 1, 1)) + i2.clip(0, 1) shouldEqual Tensor.like(i2).fromArray(Array(0, 0, 1, 1)) it("unary_-"): - (-i2) shouldEqual Tensor.fromArray(i2.shape, i2.vtype)(Array(1, 0, -1, -2)) + (-i2) shouldEqual Tensor.like(i2).fromArray(Array(1, 0, -1, -2)) describe("Boolean ops (Tensor2)"): it("inverse (!)"): - (!b2) shouldEqual Tensor2.fromArray(Axis[A], Axis[B], VType[Boolean])( + (!b2) shouldEqual Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(false, true), Array(true, false)) ) @@ -93,16 +94,16 @@ class TensorOpsElementwiseSuite extends AnyFunSpec with Matchers: it("boolean casting"): b2.asBoolean shouldEqual b2 - b2.asInt shouldEqual Tensor.fromArray(b2.shape, VType[Int])(Array(1, 0, 0, 1)) - b2.asFloat should approxEqual(Tensor.fromArray(b2.shape, VType[Float])(Array(1.0f, 0.0f, 0.0f, 1.0f))) + b2.asInt shouldEqual Tensor(b2.shape).fromArray(Array(1, 0, 0, 1)) + b2.asFloat should approxEqual(Tensor(b2.shape).fromArray(Array(1.0f, 0.0f, 0.0f, 1.0f))) it("int casting"): - i2.asBoolean shouldEqual Tensor.fromArray(i2.shape, VType[Boolean])(Array(true, false, true, true)) + i2.asBoolean shouldEqual Tensor(i2.shape).fromArray(Array(true, false, true, true)) i2.asInt shouldEqual i2 - i2.asFloat should approxEqual(Tensor.fromArray(i2.shape, VType[Float])(Array(-1.0f, 0.0f, 1.0f, 2.0f))) + i2.asFloat should approxEqual(Tensor(i2.shape).fromArray(Array(-1.0f, 0.0f, 1.0f, 2.0f))) it("float casting"): - val f2 = Tensor.fromArray(t2.shape, VType[Float])(Array(-1.1f, 0.0f, 0.9f, 2.5f)) - f2.asBoolean shouldEqual Tensor.fromArray(f2.shape, VType[Boolean])(Array(true, false, true, true)) - f2.asInt shouldEqual Tensor.fromArray(f2.shape, VType[Int])(Array(-1, 0, 0, 2)) + val f2 = Tensor.like(t2).fromArray(Array(-1.1f, 0.0f, 0.9f, 2.5f)) + f2.asBoolean shouldEqual Tensor(f2.shape).fromArray(Array(true, false, true, true)) + f2.asInt shouldEqual Tensor(f2.shape).fromArray(Array(-1, 0, 0, 2)) f2.asFloat shouldEqual f2 diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala index 55e637cf..6272a214 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala @@ -7,10 +7,10 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsFunctionalSuite extends AnyFunSpec with Matchers: - val t2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f)) ) - val t2_2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2_2 = Tensor2(Axis[A], Axis[B]).fromArray( Array(Array(10.0f, 20.0f), Array(30.0f, 40.0f)) ) @@ -22,15 +22,15 @@ class TensorOpsFunctionalSuite extends AnyFunSpec with Matchers: it("vmap over Axis A (rows)"): val res = t2.vmap(Axis[A])(_.sum) - res shouldEqual Tensor1.fromArray(Axis[A], VType[Float])(Array(3.0f, 7.0f)) + res shouldEqual Tensor1(Axis[A]).fromArray(Array(3.0f, 7.0f)) it("vmap over Axis B (columns)"): val res = t2.vmap(Axis[B])(_.sum) - res shouldEqual Tensor1.fromArray(Axis[B], VType[Float])(Array(4.0f, 6.0f)) + res shouldEqual Tensor1(Axis[B]).fromArray(Array(4.0f, 6.0f)) it("nested vmap"): val res = t2.vmap(Axis[A])(_.vmap(Axis[B])(_ => 0.0f)) - res shouldEqual Tensor.zeros(t2.shape, t2.vtype) + res shouldEqual Tensor.like(t2).fill(0.0f) describe("zipvmap (Parallel Mapping)"): @@ -38,11 +38,11 @@ class TensorOpsFunctionalSuite extends AnyFunSpec with Matchers: it("zipvmap2 adds two tensors"): val distances = zipvmap(Axis[A])(t2, t2_2)(l2) - distances should approxEqual(Tensor1.fromArray(Axis[A], VType[Float])(Array(20.12461f, 45f))) + distances should approxEqual(Tensor1(Axis[A]).fromArray(Array(20.12461f, 45f))) it("zipvmap4 adds four tensors"): val res = zipvmap(Axis[A])(t2, t2_2, t2_2, t2)((a, b, c, d) => l2(a, b) - l2(c, d)) - res should approxEqual(Tensor1.fromArray(Axis[A], VType[Float])(Array(0.0f, 0.0f))) + res should approxEqual(Tensor1(Axis[A]).fromArray(Array(0.0f, 0.0f))) describe("vapply (Axis-wise application)"): @@ -53,7 +53,7 @@ class TensorOpsFunctionalSuite extends AnyFunSpec with Matchers: it("vapply over Axis A: adds a vector to each row"): val res = t2.vapply(Axis[A])(row => row /! row.norm) - res should approxEqual(Tensor.fromArray(t2.shape, t2.vtype)( + res should approxEqual(Tensor.like(t2).fromArray( Array(0.31622776f, 0.4472136f, 0.94868326f, 0.8944272f) )) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala index 6aaa294a..bae5aef4 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsReductionSuite.scala @@ -7,22 +7,20 @@ import org.scalatest.funspec.AnyFunSpec class TensorOpsReductionSuite extends AnyFunSpec with Matchers: - val t2 = Tensor2.fromArray( + val t2 = Tensor2( Axis[A], - Axis[B], - VType[Float] - )( + Axis[B] + ).fromArray( Array( Array(1.0f, 2.0f, 3.0f), Array(4.0f, 5.0f, 6.0f) ) ) - val b2 = Tensor2.fromArray( + val b2 = Tensor2( Axis[A], - Axis[B], - VType[Boolean] - )( + Axis[B] + ).fromArray( Array( Array(true, true), Array(true, false) @@ -45,109 +43,108 @@ class TensorOpsReductionSuite extends AnyFunSpec with Matchers: it("sum axis A"): val res = t2.sum(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(5.0f, 7.0f, 9.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(5.0f, 7.0f, 9.0f))) it("sum axis B"): val res = t2.sum(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(6.0f, 15.0f))) - + res should approxEqual(Tensor.like(res).fromArray(Array(6.0f, 15.0f))) it("mean"): t2.mean shouldEqual Tensor0(3.5f) it("mean axis A"): val res = t2.mean(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(2.5f, 3.5f, 4.5f))) + res should approxEqual(Tensor.like(res).fromArray(Array(2.5f, 3.5f, 4.5f))) it("mean axis B"): val res = t2.mean(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(2.0f, 5.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(2.0f, 5.0f))) it("std"): t2.std.item should be(1.7078f +- 0.001f) it("std axis A"): val res = t2.std(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(1.5f, 1.5f, 1.5f))) + res should approxEqual(Tensor.like(res).fromArray(Array(1.5f, 1.5f, 1.5f))) it("std axis B"): val res = t2.std(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(0.8164966f, 0.8164966f))) + res should approxEqual(Tensor.like(res).fromArray(Array(0.8164966f, 0.8164966f))) it("quantile"): t2.quantile(0.5f) shouldEqual Tensor0(3.5f) it("quantile axis A"): val res = t2.quantile(0.25f, axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(1.75f, 2.75f, 3.75f))) + res should approxEqual(Tensor.like(res).fromArray(Array(1.75f, 2.75f, 3.75f))) it("quantile axis B"): val res = t2.quantile(0.25f, axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(1.5f, 4.5f))) + res should approxEqual(Tensor.like(res).fromArray(Array(1.5f, 4.5f))) it("median"): t2.median shouldEqual Tensor0(3.5f) it("median axis A"): val res = t2.median(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(2.5f, 3.5f, 4.5f))) + res should approxEqual(Tensor.like(res).fromArray(Array(2.5f, 3.5f, 4.5f))) it("median axis B"): val res = t2.median(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(2.0f, 5.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(2.0f, 5.0f))) it("max"): t2.max shouldEqual Tensor0(6.0f) it("max axis A"): val res = t2.max(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(4.0f, 5.0f, 6.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(4.0f, 5.0f, 6.0f))) it("max axis B"): val res = t2.max(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(3.0f, 6.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(3.0f, 6.0f))) it("min"): t2.min shouldEqual Tensor0(1.0f) it("min axis A"): val res = t2.min(axis = Axis[A]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(1.0f, 2.0f, 3.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(1.0f, 2.0f, 3.0f))) it("min axis B"): val res = t2.min(axis = Axis[B]) - res should approxEqual(Tensor.fromArray(res.shape, res.vtype)(Array(1.0f, 4.0f))) + res should approxEqual(Tensor.like(res).fromArray(Array(1.0f, 4.0f))) it("argmax"): t2.argmax shouldEqual Tensor0(5) it("argmax axis A"): val res = t2.argmax(axis = Axis[A]) - res shouldEqual Tensor.fromArray(res.shape, res.vtype)(Array(1, 1, 1)) + res shouldEqual Tensor.like(res).fromArray(Array(1, 1, 1)) it("argmax axis B"): val res = t2.argmax(axis = Axis[B]) - res shouldEqual Tensor.fromArray(res.shape, res.vtype)(Array(2, 2)) + res shouldEqual Tensor.like(res).fromArray(Array(2, 2)) it("argmin"): t2.argmin shouldEqual Tensor0(0) it("argmin axis A"): val res = t2.argmin(axis = Axis[A]) - res shouldEqual Tensor.fromArray(res.shape, res.vtype)(Array(0, 0, 0)) + res shouldEqual Tensor.like(res).fromArray(Array(0, 0, 0)) it("argmin axis B"): val res = t2.argmin(axis = Axis[B]) - res shouldEqual Tensor.fromArray(res.shape, res.vtype)(Array(0, 0)) + res shouldEqual Tensor.like(res).fromArray(Array(0, 0)) describe("Boolean Reductions"): it("all"): b2.all shouldEqual Tensor0(false) - val allTrue = Tensor.ones(b2.shape, b2.vtype) + val allTrue = Tensor.like(b2).fill(true) allTrue.all shouldEqual Tensor0(true) it("any"): b2.any shouldEqual Tensor0(true) - val allFalse = Tensor.zeros(b2.shape, b2.vtype) + val allFalse = Tensor.like(b2).fill(false) allFalse.any shouldEqual Tensor0(false) describe("Approximate Equality") { diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala index 7175d571..e9b175ca 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsStructureSuite.scala @@ -10,7 +10,7 @@ import scala.compiletime.testing.typeCheckErrors class TensorOpsStructureSuite extends AnyFunSpec with Matchers: // Shape: A=2, B=2, C=1 - val t3 = Tensor3.fromArray(Axis[A], Axis[B], Axis[C], VType[Float])( + val t3 = Tensor3(Axis[A], Axis[B], Axis[C]).fromArray( Array( Array(Array(1.0f), Array(2.0f)), Array(Array(3.0f), Array(4.0f)) @@ -33,7 +33,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: val res = t3.rearrange((Axis[A |*| B], Axis[C])) res.axes shouldBe List("A*B", "C") res should approxEqual( - Tensor2.fromArray(Axis[A |*| B], Axis[C], VType[Float])( + Tensor2(Axis[A |*| B], Axis[C]).fromArray( Array(Array(1.0f), Array(2.0f), Array(3.0f), Array(4.0f)) ) ) @@ -61,7 +61,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: "flattened.rearrange((Axis[A], Axis[B], Axis[C]), Axis[A] -> 2, Axis[B] -> 2)" should compile it("complex reshape fails incrementally: (a b) (c d) -> (a c) (b d)"): - val t = Tensor2.fromArray(Axis[A |*| B], Axis[C |*| D], VType[Float])( + val t = Tensor2(Axis[A |*| B], Axis[C |*| D]).fromArray( Array( Array(1.0f, 2.0f, 3.0f, 4.0f), Array(5.0f, 6.0f, 7.0f, 8.0f), @@ -93,7 +93,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: "t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[C] -> 2, Axis[B] -> 2, Axis[D] -> 2)" should compile it("complex rearrange: (a b) (c d) -> (a c) (b d)"): - val t = Tensor2.fromArray(Axis[A |*| B], Axis[C |*| D], VType[Float])( + val t = Tensor2(Axis[A |*| B], Axis[C |*| D]).fromArray( Array( Array(1.0f, 2.0f, 3.0f, 4.0f), Array(5.0f, 6.0f, 7.0f, 8.0f), @@ -104,7 +104,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: t.axes shouldBe List("A*B", "C*D") val res = t.rearrange((Axis[A |*| C], Axis[B |*| D]), Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2) res.axes shouldBe List("A*C", "B*D") - res should approxEqual(Tensor2.fromArray(Axis[A |*| C], Axis[B |*| D], VType[Float])( + res should approxEqual(Tensor2(Axis[A |*| C], Axis[B |*| D]).fromArray( Array( Array(1.0f, 2.0f, 5.0f, 6.0f), Array(3.0f, 4.0f, 7.0f, 8.0f), @@ -124,14 +124,14 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: describe("split function"): it("split axis into two axes"): - val ab = Tensor.ones(Shape(Axis[A] -> 4, Axis[B] -> 12), VType[Float]) + val ab = Tensor(Shape(Axis[A] -> 4, Axis[B] -> 12)).fill(1f) val acd = ab.split(Axis[B], Axis[C] -> 6, Axis[D] -> 2) acd.axes shouldBe (List("A", "C", "D")) acd.shape(Axis[C]) shouldBe (6) acd.shape(Axis[D]) shouldBe (2) it("split axis into two axes, one being the type of original axis"): - val ab = Tensor.ones(Shape(Axis[A] -> 4, Axis[B] -> 12), VType[Float]) + val ab = Tensor(Shape(Axis[A] -> 4, Axis[B] -> 12)).fill(1f) val acd = ab.split(Axis[B], Axis[B] -> 6, Axis[D] -> 2) acd.axes shouldBe (List("A", "B", "D")) acd.shape(Axis[B]) shouldBe (6) @@ -165,13 +165,13 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: t3.relabel(Axis[C] -> Axis[X]).axes shouldBe List("A", "B", "X") it("relabel all axes"): - val t = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array.fill(2, 2)(1.0f)) + val t = Tensor2(Axis[A], Axis[B]).fromArray(Array.fill(2, 2)(1.0f)) val relabeled = t.relabelAll((Axis[C], Axis[D])) relabeled.axes shouldBe List("C", "D") describe("tril / triu"): - val t = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(1.0f, 2.0f), Array(3.0f, 4.0f) @@ -194,13 +194,13 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: describe("where"): - val t1 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t1 = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(1.0f, 2.0f), Array(3.0f, 4.0f) ) ) - val t2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + val t2 = Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(10.0f, 20.0f), Array(30.0f, 40.0f) @@ -208,14 +208,14 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: ) it("uniform mask"): - val mask = Tensor.zeros(t1.shape, VType[Boolean]) + val mask = Tensor(t1.shape).fill(false) where(mask, t1, t2) should approxEqual(t2) where(!mask, t1, t2) should approxEqual(t1) it("triu mask"): - val mask = triu(Tensor.ones(t1.shape, VType[Boolean])) + val mask = triu(Tensor(t1.shape).fill(true)) where(mask, t1, t2) should approxEqual( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( + Tensor2(Axis[A], Axis[B]).fromArray( Array( Array(1.0f, 2.0f), Array(30.0f, 4.0f) @@ -227,7 +227,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: it("Prime axes are rearrangable"): // As rearrange uses einops the "+" om the derived label for B |+| C must be handled in the rearrange operation to not trigger error - val t = Tensor2.fromArray(Axis[A], Axis[Prime[B] |*| B], VType[Float])( + val t = Tensor2(Axis[A], Axis[Prime[B] |*| B]).fromArray( Array(Array(1.0f, 2.0f, 3.0f, 4.0f), Array(5.0f, 6.0f, 7.0f, 8.0f)) ) val tRearranged = t.rearrange( @@ -238,15 +238,15 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: it("|+| axes are rearrangable"): // As rearrange uses einops the "+" om the derived label for B |+| C must be handled in the rearrange operation to not trigger error - val t = Tensor2.fromArray(Axis[A], Axis[B |+| C], VType[Float])( + val t = Tensor2(Axis[A], Axis[B |+| C]).fromArray( Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f)) ) val tRearranged = t.rearrange((Axis[B |+| C], Axis[A])) tRearranged.axes shouldBe List("B+C", "A") it("concatenate2 same axes"): - val part1 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(1.0f, 2.0f))) - val part2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(3.0f, 4.0f))) + val part1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) + val part2 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 4.0f))) val joined = concatenate(part1, part2, Axis[B]) joined.axes shouldBe List("A", "B") joined.shape(Axis[B]) shouldBe (part1.shape(Axis[B]) + part2.shape(Axis[B])) @@ -254,9 +254,9 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: joined.slice(Axis[B] -> (part1.shape(Axis[B]) until (part1.shape(Axis[B]) + part2.shape(Axis[B])))) should approxEqual(part2) it("concatenateN same axes"): - val part1 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(1.0f, 2.0f))) - val part2 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(3.0f, 4.0f))) - val part3 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(3.0f, 4.0f))) + val part1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) + val part2 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 4.0f))) + val part3 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(3.0f, 4.0f))) val joined = concatenate(Seq(part1, part2, part3), Axis[B]) joined.axes shouldBe List("A", "B") joined.shape(Axis[B]) shouldBe (part1.shape(Axis[B]) + part2.shape(Axis[B]) + part3.shape(Axis[B])) @@ -265,8 +265,8 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: joined.slice(Axis[B] -> ((part1.shape(Axis[B]) + part2.shape(Axis[B])) until (part1.shape(Axis[B]) + part2.shape(Axis[B]) + part3.shape(Axis[B])))) should approxEqual(part3) it("concatenate2 different axes"): - val part1 = Tensor2.fromArray(Axis[A], Axis[B], VType[Float])(Array(Array(1.0f, 2.0f))) - val part2 = Tensor2.fromArray(Axis[A], Axis[C], VType[Float])(Array(Array(3.0f, 4.0f))) + val part1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) + val part2 = Tensor2(Axis[A], Axis[C]).fromArray(Array(Array(3.0f, 4.0f))) val joined = concatenate(part1, part2) joined.axes shouldBe List("A", "B+C") joined.shape(Axis[B |+| C]) shouldBe (part1.shape(Axis[B]) + part2.shape(Axis[C])) @@ -276,7 +276,7 @@ class TensorOpsStructureSuite extends AnyFunSpec with Matchers: describe("Deconcatenation"): it("deconcatenate on |+| axis"): - val t = Tensor2.fromArray(Axis[A], Axis[B |+| C], VType[Float])( + val t = Tensor2(Axis[A], Axis[B |+| C]).fromArray( Array(Array(1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f)) ) val (partB, partC) = t.deconcatenate(Axis[B |+| C], (Axis[B] -> 2, Axis[C] -> 3)) diff --git a/core/src/test/scala/dimwit/tensor/TensorWithValueClassSuite.scala b/core/src/test/scala/dimwit/tensor/TensorWithValueClassSuite.scala index fe714d53..96e80396 100644 --- a/core/src/test/scala/dimwit/tensor/TensorWithValueClassSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorWithValueClassSuite.scala @@ -21,7 +21,7 @@ class TensorWithValueClassSuite extends AnyFunSpec with Matchers: given IsFloat[V2] with {} // make all IsFloat ops available import ValueClassScope.* - val t = Tensor.zeros(Shape(Axis[A] -> 1, Axis[B] -> 2), VType[Float]) + val t = Tensor(Shape(Axis[A] -> 1, Axis[B] -> 2)).fill(0f) val v1 = V1(t) val v2 = V2(t) "v1 + v1" should compile diff --git a/docs/README.md b/docs/README.md index b7265cc3..a6f3f148 100644 --- a/docs/README.md +++ b/docs/README.md @@ -34,9 +34,9 @@ trait Batch derives Label trait Feature derives Label // Create a 2D tensor with shape (3, 2), labeled with Batch and Feature -val t = Tensor.fromArray( +val t = Tensor( Shape(Axis[Batch] -> 3, Axis[Feature] -> 2), - VType[Float])( +).fromArray( Array( 1.0f, 2.0f, 3.0f, 4.0f, diff --git a/examples/src/main/scala/api/AutodiffAPI.scala b/examples/src/main/scala/api/AutodiffAPI.scala deleted file mode 100644 index 4c2e48f0..00000000 --- a/examples/src/main/scala/api/AutodiffAPI.scala +++ /dev/null @@ -1,151 +0,0 @@ -package examples.api - -import dimwit.* - -private trait A derives Label -private trait B derives Label -private trait C derives Label -private trait D derives Label - -@main -def autoDiffAPI(): Unit = - val AB = Tensor.ones( - Shape( - Axis[A] -> 10, - Axis[B] -> 5 - ), - VType[Float] - ) - val AC = Tensor.ones( - Shape( - Axis[A] -> 10, - Axis[C] -> 5 - ), - VType[Float] - ) - val ABCD = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4, - Axis[D] -> 5 - ), - VType[Float] - ) - { - def f(x: Tensor1[A, Float]): Tensor0[Float] = x.sum - val df = Autodiff.grad(f) - val delta = df(Tensor1.fromArray(Axis[A], VType[Float])(Array.fill(10)(1.0f))) - println(delta.value.shape) // Access underlying value - } - { - type ParamsTuple = (Tensor2[A, B, Float], Tensor1[C, Float]) - def f(params: ParamsTuple): Tensor0[Float] = - params._1.sum + params._2.sum - val df = Autodiff.grad(f) - val delta = df( - ( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( - Array( - Array.fill(5)(1.0f), - Array.fill(5)(1.0f) - ) - ), - Tensor1.fromArray(Axis[C], VType[Float])(Array.fill(5)(1.0f)) - ) - ) - println((delta.value._1.shape, delta.value._2.shape)) - } - { - case class Params( - a: Tensor2[A, B, Float], - b: Tensor1[C, Float] - ) derives TensorTree - def f(params: Params): Tensor0[Float] = - params.a.sum + params.b.sum - val df = Autodiff.grad(f) - val delta = df( - Params( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( - Array( - Array.fill(5)(1.0f), - Array.fill(5)(1.0f) - ) - ), - Tensor1.fromArray(Axis[C], VType[Float])(Array.fill(5)(1.0f)) - ) - ) - println(delta) - } - { - def f(x: Tensor1[A, Float]): Tensor1[A, Float] = x - val df = Autodiff.jacobian(f) - val delta = df(Tensor1.fromArray(Axis[A], VType[Float])(Array.fill(10)(1.0f))) - println(delta.shape) - } - { - def f(x: Tensor1[A, Float]) = x.outerProduct(x) - val df = Autodiff.jacobian(f) - val delta = df(Tensor1.fromArray(Axis[A], VType[Float])(Array.fill(10)(1.0f))) - println(delta.shape) - } - { - import dimwit.tensor.TensorOps.* - type ParamsTuple = (Tensor2[A, B, Float], Tensor1[C, Float]) - def f(x: ParamsTuple): Tensor1[A, Float] = x._1.slice(Axis[B] -> 0) - val df = Autodiff.jacobian(f) - val delta = df( - ( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( - Array( - Array.fill(5)(1.0f), - Array.fill(5)(1.0f) - ) - ), - Tensor1.fromArray(Axis[C], VType[Float])(Array.fill(5)(1.0f)) - ) - ) - println((delta._1.shape, delta._2.shape)) - } - { - println("Hessian") - def f(x: Tensor1[A, Float]): Tensor0[Float] = x.sum - val df = Autodiff.jacobian(f) - val ddf = Autodiff.jacobian(df) - val delta = ddf(Tensor1.fromArray(Axis[A], VType[Float])(Array.fill(10)(1.0f))) - println(delta.shape) - } - { - def f(x: Tensor1[A, Float]): Tensor0[Float] = x.sum - val df = Autodiff.jacobian(f) - val ddf = Autodiff.jacobian(df) - val dddf = Autodiff.jacobian(ddf) - val ddddf = Autodiff.jacobian(dddf) - val delta = ddddf(Tensor1.fromArray(Axis[A], VType[Float])(Array.fill(10)(1.0f))) - println(delta.shape) - } - { - import dimwit.tensor.TensorOps.* - type ParamsTuple = (Tensor2[A, B, Float], Tensor1[C, Float]) - def f(x: ParamsTuple): Tensor0[Float] = x._1.sum - val df = Autodiff.jacobian(f) - val ddf = Autodiff.jacobian(df) - val delta = ddf( - ( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( - Array( - Array.fill(5)(1.0f), - Array.fill(5)(1.0f) - ) - ), - Tensor1.fromArray(Axis[C], VType[Float])(Array.fill(5)(1.0f)) - ) - ) - // TODO Is this actually correct, check it! - println( - ( - (delta._1._1.shape, delta._1._2.shape), - (delta._2._1.shape, delta._2._2.shape) - ) - ) - } diff --git a/examples/src/main/scala/api/TensorAPI.scala b/examples/src/main/scala/api/TensorAPI.scala deleted file mode 100644 index f084ef6e..00000000 --- a/examples/src/main/scala/api/TensorAPI.scala +++ /dev/null @@ -1,621 +0,0 @@ -package examples.api - -import dimwit.* -import dimwit.Conversions.given -import me.shadaj.scalapy.py -import me.shadaj.scalapy.py.PythonException - -def opBlock[T](operation: String)(block: => T): Unit = - val res = block - block match - case t: Tensor[?, Float] => - println(f"$operation%-30s: ${t.shape}%-30s == ${py.eval("res.shape if hasattr(res, 'shape') else res")}") - case v => - println(f"$operation%-30s: $v%-30s == ${py.eval("res")}") - -@main -def tensorAPI(): Unit = - - trait A derives Label - trait A1 derives Label - trait A2 derives Label - trait B derives Label - trait C derives Label - trait D derives Label - trait E derives Label - trait X derives Label - trait Y derives Label - - trait Stack derives Label - - val UNSUPPORTED = "Not Supported by dimwit" - py.exec("import jax") - py.exec("import jax.numpy as jnp") - py.exec("import einops") - // py.eval("import jax.numpy as jnp") - val AB = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - py.exec("ab = jnp.ones((2, 3))") - val AC = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[C] -> 4 - ), - VType[Float] - ) - py.exec("ac = jnp.ones((2, 4))") - val BCD = Tensor.ones( - Shape( - Axis[B] -> 3, - Axis[C] -> 4, - Axis[D] -> 5 - ), - VType[Float] - ) - py.exec("bcd = jnp.ones((3, 4, 5))") - val BC = Tensor.ones( - Shape( - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - py.exec("bc = jnp.ones((3, 4))") - val ABCD = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4, - Axis[D] -> 5 - ), - VType[Float] - ) - py.exec("abcd = jnp.ones((2, 3, 4, 5))") - println((AB.shape, py.eval("ab.shape"))) - println((AC.shape, py.eval("ac.shape"))) - println((ABCD.shape, py.eval("abcd.shape"))) - { - /* - * BROADCASTING - * https://numpy.org/doc/stable/user/basics.broadcasting.html - */ - println("BROADCASTING") - opBlock("Axes broadcasting backward: ABCD + BCD") { - py.exec("res = abcd + bcd") - BCD +! ABCD - ABCD +! BCD - } - opBlock("Scalar broadcast: ABCD + Scalar") { - py.exec("res = abcd + 5") - ABCD +! 5f - } - opBlock("Axes broadcasting backward: ABCD + CD") { - py.exec("cd = jnp.ones((4,5))") - py.exec("res = abcd + cd") - val CD = Tensor.ones( - Shape( - Axis[C] -> 4, - Axis[D] -> 5 - ), - VType[Float] - ) - ABCD +! CD - } - opBlock("Dim broadcast: ABC1 to ABCD") { - // - py.exec("abc1 = jnp.ones((2,3,4,1))") - py.exec("res = abcd + abc1") - UNSUPPORTED // Do not support as d != 1, broadcasting does implicit magic - } - opBlock("Dims broadcast: AB11 to ABCD") { - py.exec("ab11 = jnp.ones((2,3,1,1))") - py.exec("res = abcd + ab11") - ABCD - } - opBlock("Magic broadcast: a1 + b") { - py.exec("a1 = jnp.ones((2,1))") - py.exec("b = jnp.ones((3))") - py.exec("res = a1 + b") - AB - } - // Negative examples - opBlock("Axes broadcasting forward: ABCD + AB") { // Axes broadcasting (forward) - try - py.exec("res = abcd + ab") - assert(false, "Expected exception not thrown") - catch - case e: PythonException => - py.exec("res = 'Not Supported by JAX'") - ABCD +! AB - } - opBlock("Axes broadcasting forward and backward: : ABCD + BC") { // Axes broadcasting (forward and backward) - try - py.exec("res = abcd + bc") - assert(false, "Expected exception not thrown") - catch - case e: PythonException => - py.exec("res = 'Not Supported by JAX'") - ABCD +! BC - } - - /** ELEMENT-WISE OPERATIONS - */ - println("ELEMENT-WISE OPERATIONS") - opBlock("+") { - py.exec("res = ab + ab") - AB + AB - } - opBlock("*") { - py.exec("res = ab * ab") - AB * AB - } - opBlock("-") { - py.exec("res = ab - ab") - AB - AB - } - opBlock("/") { - py.exec("res = ab / ab") - AB / AB - } - opBlock("abs") { - py.exec("res = jnp.abs(ab)") - AB.abs - } - opBlock("sign") { - py.exec("res = jnp.sign(ab)") - AB.sign - } - opBlock("pow") { - py.exec("res = ab ** 2") - AB.pow(2) - } - opBlock("sqrt") { - py.exec("res = jnp.sqrt(ab)") - AB.sqrt - } - opBlock("exp") { - py.exec("res = jnp.exp(ab)") - AB.exp - } - opBlock("log") { - py.exec("res = jnp.log(ab)") - AB.log - } - opBlock("sin") { - py.exec("res = jnp.sin(ab)") - AB.sin - } - opBlock("cos") { - py.exec("res = jnp.cos(ab)") - AB.cos - } - opBlock("tanh") { - py.exec("res = jnp.tanh(ab)") - AB.tanh - } - opBlock("clip") { - py.exec("res = jnp.clip(ab, 0, 1)") - AB.clip(0, 1) - } - opBlock("<") { - py.exec("res = ab < ab") - AB < AB - } - opBlock(">") { - py.exec("res = ab > ab") - AB > AB - } - opBlock("<=") { - py.exec("res = ab <= ab") - AB <= AB - } - opBlock(">=") { - py.exec("res = ab >= ab") - AB >= AB - } - opBlock("==") { - py.exec("res = jnp.array_equal(ab, ab)") - AB == AB - } - - /** REDUCTION - */ - opBlock("sum") { - py.exec("res = jnp.sum(ab)") - AB.sum - } - opBlock("sum ab axis=0") { - py.exec("res = jnp.sum(ab, axis=0)") - AB.sum(Axis[A]) - } - opBlock("sum abcd axis=0") { - py.exec("res = jnp.sum(abcd, axis=0)") - ABCD.sum((Axis[A])) - } - opBlock("sum ab axis=(0,1)") { - py.exec("res = jnp.sum(ab, axis=(0,1))") - AB.sum((Axis[A], Axis[B])) - } - opBlock("sum abcd axis=(0,1)") { - py.exec("res = jnp.sum(abcd, axis=(0,1))") - ABCD.sum((Axis[A], Axis[B])) - } - opBlock("mean") { - py.exec("res = jnp.mean(ab)") - AB.mean - } - opBlock("mean ab axis=0") { - py.exec("res = jnp.mean(ab, axis=0)") - AB.mean(Axis[A]) - } - opBlock("max ab") { - py.exec("res = jnp.max(ab)") - AB.max - } - opBlock("max ab axis=0") { - py.exec("res = jnp.max(ab, axis=0)") - AB.max(Axis[A]) - } - opBlock("min ab") { - py.exec("res = jnp.min(ab)") - AB.min - } - opBlock("min ab axis=0") { - py.exec("res = jnp.min(ab, axis=0)") - AB.min(Axis[A]) - } - opBlock("argmax ab") { - py.exec("res = jnp.argmax(ab)") - AB.argmax - } - opBlock("argmax ab axis=0") { - py.exec("res = jnp.argmax(ab, axis=0)") - AB.argmax(Axis[A]) - } - opBlock("argmin ab") { - py.exec("res = jnp.argmin(ab)") - AB.argmin - } - opBlock("argmin ab axis=0") { - py.exec("res = jnp.argmin(ab, axis=0)") - AB.argmin(Axis[A]) - } - - /** CONTRACT Analog to JAX tensordot with a single axis, with two changes: - * - Only a single axis is allowed TODO allow multiple axes - */ - opBlock("dot") { - py.exec("res = einops.einsum(ab, ac, 'a b, a c -> b c')") // einsum variant - py.exec("res = jnp.tensordot(ab, ac, axes=(0, 0))") // pure JAX variant - AB.dot(Axis[A])(AC) - } - opBlock("dot abcd axis=2") { - py.exec("res = einops.einsum(abcd, abcd, 'a b c d, e f c g -> a b d e f g')") // einsum variant - py.exec("res = jnp.tensordot(abcd, abcd, axes=(2, 2))") // pure JAX variant - ABCD.dot(Axis[C])(ABCD) - } - - /** OUTER PRODUCT (contract over zero axes) Analog to JAX outer product, i.e., no axes to contract - */ - opBlock("outerProduct") { - py.exec("res = jnp.einsum('ij, kl -> ijkl', ab, ac)") // einsum variant - py.exec("res = jnp.tensordot(ab, ac, axes=0)") // pure JAX variant - AB.outerProduct(AC) - } - opBlock("stack AB AB AB AB") { - py.exec("res = jnp.stack([ab, ab, ab, ab], axis=0)") - stack(List(AB, AB, AB, AB), Axis[Stack]) - } - opBlock("stack AB AB AB AB axis=1") { - py.exec("res = jnp.stack([ab, ab, ab, ab], axis=1)") - stack(List(AB, AB, AB, AB), Axis[Stack], afterAxis = Axis[A]) - } - opBlock("concat AB AB") { - py.exec("res = jnp.concatenate([ab, ab], axis=0)") - concatenate(List(AB, AB), Axis[A]) - } - opBlock("concat AB AB axis=1") { - py.exec("res = jnp.concatenate([ab, ab], axis=1)") - concatenate(List(AB, AB), Axis[B]) - } - - /** SLICE Analog to JAX slice(...) or JAX at(...).get, with two changes: - * - Out of range index leads to an error (instead of clipping) - * - No colon access (e.g., X[:, 0]), as due to name of axes this is not necessary (just leave out name) - */ - // Select single index - opBlock("slice ab axis=0") { - py.exec("res = ab[0, :]") - AB.slice(Axis[A] -> 0) - } - opBlock("slice abcd axis=3") { - py.exec("res = abcd[:, :, :, 0]") - ABCD.slice(Axis[D] -> 0) - } - opBlock("slice ab axis=0:1") { - py.exec("res = ab[0:1, :]") - AB.slice(Axis[A] -> (0 until 1)) - } - opBlock("slice ab axis=0,2") { - py.exec("res = ab[0:1, 2]") - AB.slice( - ( // TODO make (()) optional - Axis[A] -> (0 until 1), - Axis[B] -> 2 - ) - ) - } - // opBlock("slice ab axis=[0,2]") { // TODO make this work - // py.exec("res = ab[:, [0,2]]") - // AB.slice( // TODO make (()) optional - // Axis[B] -> List(0, 2), - // ) - // } - opBlock("gather ab axis=1") { - py.exec("indices = jnp.array([0,2])") - py.exec("res = jnp.take(ab, indices, axis=1)") - val res = AB.take(Axis[B])( - Tensor1.fromArray( - Axis[C], - VType[Int] - )( - Array(0, 2) - ) - ) - res - } - - /** SET Analog to JAX at(...).set, with two changes: - * - Out of range index leads to an error (instead of clipping) - * - No colon access (e.g., X[:, 0]), as due to name of axes this is not necessary (just leave out name) - */ - opBlock("set ab axis=0") { - py.exec("res = ab.at[0, :].set(jnp.array([0,1,2]))") - AB.set( - Axis[A] -> 0 - )(Tensor1.fromArray(Axis[B], VType[Float])(Array(0, 1, 2))) - } - // set sub-matrix, AB.at[0:1, 0:1].set([[1,2],[3,4]]) - opBlock("set ab axis=0:1,0:1") { - py.exec("res = ab.at[0:2, 0:2].set(jnp.array([[1,2],[3,4]]))") - AB.set( - ( // TODO make (()) optional - Axis[A] -> (0 until 2), - Axis[B] -> (0 until 2) - ) - )( - Tensor2.fromArray(Axis[A], Axis[B], VType[Float])( - Array( - Array(1f, 2f), - Array(3f, 4f) - ) - ) - ) - } - - /** REARRANGE Analog to einops rearrange, but with named axes. For JAX this replaces `transpose` and `reshape` operations. - */ - // einops.rearrange(ABCD, 'a b c d -> b a c d') - opBlock("rearrange ABCD swap A and B") { - py.exec("res = einops.rearrange(abcd, 'a b c d -> b a c d')") // einops variant - py.exec("res = jnp.transpose(abcd, (1, 0, 2, 3))") // pure JAX variant - // TODO maybe rename to `transpose` as in JAX? - ABCD.rearrange( - (Axis[B], Axis[A], Axis[C], Axis[D]) - ) - } - opBlock("rearrange ABCD flatten A and B") { - py.exec("res = einops.rearrange(abcd, 'a b c d -> (b a) c d')") // einops variant - py.exec("res = jnp.reshape(abcd.transpose((1, 0, 2, 3)), (abcd.shape[0]*abcd.shape[1], abcd.shape[2], abcd.shape[3]))") // pure JAX variant - ABCD.rearrange( - (Axis[B |*| A], Axis[C], Axis[D]) - ) - } - opBlock("rearrange ABCD unflatten AB") { - py.exec("tmp = einops.rearrange(abcd, 'a b c d -> (b a) c d')") // Setup - py.exec("res = einops.rearrange(tmp, '(b a) c d -> a b c d', a=abcd.shape[0], b=abcd.shape[1])") // einops variant - py.exec("res = jnp.reshape(tmp, (abcd.shape[0], abcd.shape[1], tmp.shape[1], tmp.shape[2])).transpose((1, 0, 2, 3))") // pure JAX variant - val ABCDFlat = ABCD.rearrange( - (Axis[B |*| A], Axis[C], Axis[D]) - ) - ABCDFlat.rearrange( - (Axis[A], Axis[B], Axis[C], Axis[D]), - (ABCD.shape.dim(Axis[A]), ABCD.shape.dim(Axis[B])) - ) - } - opBlock("chunk ABCD") { - py.exec("res = list(map(lambda x: x.shape, jnp.array_split(abcd, 2, axis=2)))") - ABCD.chunk(Axis[C], 2).map(_.shape) - } - - /** AS / RELABEL - rename axes labels */ - opBlock("as AB to XY") { - py.exec("res = ab # no equivalent in JAX, as axes are not named") - AB.relabel(Axis[A] -> Axis[X]).relabel(Axis[B] -> Axis[Y]) - // TODO add relabel of tuple of axes? => AB.relabel((Axis[A] -> Axis[X], Axis[B] -> Axis[Y])) - } - opBlock("relabel AB to XB") { - py.exec("res = ab # no equivalent in JAX, as axes are not named") - AB.relabel( - Axis[A] -> Axis[X] - ) - } - - /** SWAP */ - opBlock("swap AB axes A and B") { - py.exec("res = jnp.swapaxes(ab, 0, 1)") - AB.swap(Axis[A], Axis[B]) - } - - /** RAVEL */ - // AB.ravel() - opBlock("ravel AB") { - py.exec("res = ab.ravel()") - val res = ABCD.ravel - res - } - - /** APPEND AXIS Analog to jnp.expand_dims / None indexing in JAX, adds a new axis at the end or beginning. If axis must be inserted at a specific position use `rearrange` after `appendAxis` or `prependAxis`. - */ - // AB[:, :, None] - opBlock("append axis C to AB") { - py.exec("res = ab[:, :, None]") - AB.appendAxis(Axis[C]) - } - opBlock("prepend axis C to AB") { - py.exec("res = ab[None, :, :]") - AB.prependAxis(Axis[C]) - } - opBlock("insert axis C to AB") { - py.exec("res = ab[:, None, :]") - /* - Note that we have no direct equivalent to this in dimwit, - but we can achieve the same result by first appending or prepending the axis, - and then rearranging the axes to the desired order. - */ - AB.prependAxis(Axis[C]) - .rearrange( - (Axis[A], Axis[C], Axis[B]) - ) - } - - /** SQUEEZE Analog to jnp.squeeze in JAX - */ - opBlock("squeeze A from AB") { - py.exec("tmp = jnp.ones((1,3))") // Setup - val tmp = Tensor.ones( - Shape( - Axis[A] -> 1, - Axis[B] -> 3 - ), - VType[Float] - ) - py.exec("res = jnp.squeeze(tmp, axis=0)") - tmp.squeeze(Axis[A]) - } - - /** VMAP (/ ZIPVMAP) Analog to JAX vmap, with one changes: - * - vmap allows only single axis - * - zipvmap for multiple tensors to be mapped over the same axis (vmap in JAX) - */ - opBlock("vmap AB over axis A") { - py.exec("res = jax.vmap(lambda row: jnp.sum(row))(ab)") - AB.vmap(Axis[A]) { row => row.sum } - } - opBlock("vmap ABCD over axis C") { - py.exec("res = jax.vmap(lambda slice: jnp.sum(slice, axis=0), in_axes=2)(abcd)") - ABCD.vmap(Axis[C]) { slice => slice.sum(Axis[A]) } - } - opBlock("vmap ABCD over axis C and D") { - // TODO Is this even supported in JAX? - // py.exec("res = jax.vmap(lambda ad: jnp.sum(ad, axis=0), in_axes=(1, 2))(abcd)") - // TODO ABCD.vmap((Axis[B], Axis[D])) { _.sum(Axis[A]) } - ABCD - } - opBlock("vmap ABCD over axis C then D") { - py.exec("res = jax.vmap(lambda abc: jax.vmap(lambda ad: jnp.sum(ad, axis=0), in_axes=1)(abc), in_axes=2)(abcd)") - ABCD.vmap(Axis[C]) { xx => - xx.vmap(Axis[B]) { x => - val res = x.sum(Axis[A]) - res - } - } - } - opBlock("vmap/zipvmap AB AC") { - py.exec("res = jax.vmap(lambda abi, aci: jnp.sum(abi) + jnp.sum(aci))(ab, ac)") - zipvmap(Axis[A])((AB, AC)) { case (abi, aci) => - abi.sum + aci.sum - } - } - opBlock("vmap/zipvmap AB AC AB AC") { - py.exec("res = jax.vmap(lambda abi, aci, ab2i, ac2i: jnp.sum(abi) + jnp.sum(aci) + jnp.sum(ab2i) + jnp.sum(ac2i))(ab, ac, ab, ac)") - zipvmap(Axis[A])((AB, AC, AB, AC)) { case (abi, aci, ab2i, ac2i) => - abi.sum + aci.sum + ab2i.sum + ac2i.sum - } - } - opBlock("vapply AB over axis A") { - py.exec("res = jnp.apply_along_axis(lambda row: row, 0, ab)") - val res = AB.vapply(Axis[A]) { row => row } - res - } - - /** WHERE Analog to jnp.where in JAX - */ - opBlock("where") { - py.exec("shape = (2, 3)") - py.exec("x = jnp.ones(shape)") - py.exec("y = jnp.zeros(shape)") - py.exec("condition = jnp.zeros(shape)") - py.exec("res = jnp.where(condition, x, y)") - val shape = Shape(Axis[A] -> 2, Axis[B] -> 3) - val x = Tensor.ones(shape, VType[Float]) - val y = Tensor.zeros(shape, VType[Float]) - val condition = Tensor.zeros(shape, VType[Boolean]) - where(condition, x, y) - } - } - { - - /** LINEAR ALGEBRA - */ - opBlock("trace ABCD over axes A and B") { - py.exec("res = jnp.trace(abcd, axis1=0, axis2=1)") - ABCD.trace(Axis[A], Axis[B]) - } - opBlock("diagonal AB") { - py.exec("res = jnp.diagonal(ab)") - AB.diagonal - } - opBlock("diagonal ABCD over axes A and B") { - py.exec("res = jnp.diagonal(abcd, axis1=0, axis2=1)") - ABCD.diagonal(Axis[A], Axis[B]) - } - opBlock("trace AB") { - py.exec("res = jnp.trace(ab)") - AB.trace - } - opBlock("det abc1c2 over axes C1 and C2") { - trait C1 derives Label - trait C2 derives Label - val ABC1C2 = Tensor.ones( - Shape(Axis[A] -> 2, Axis[B] -> 3, Axis[C1] -> 4, Axis[C2] -> 4), - VType[Float] - ) - py.exec("abc1c2 = jnp.ones((2, 3, 4, 4))") - py.exec("res = jnp.linalg.det(abc1c2)") - ABC1C2.det(Axis[C1], Axis[C2]) - } - opBlock("det A1A2") { - val A1A2 = Tensor.ones( - Shape(Axis[A1] -> 2, Axis[A2] -> 2), - VType[Float] - ) - py.exec("a1a2 = jnp.ones((2, 2))") - py.exec("res = jnp.linalg.det(a1a2)") - A1A2.det - } - opBlock("norm AB") { - py.exec("res = jnp.linalg.norm(ab)") - AB.norm - } - opBlock("inv AB") { - trait A1 derives Label - trait A2 derives Label - val a1a2 = Tensor2 - .fromArray(Axis[A1], Axis[A2], VType[Float])( - Array( - Array(2f, 0f), - Array(0f, 2f) - ) - ) - .toDevice(Device.CPU) - py.exec("a1a2 = jnp.array([[2, 0],[0, 2]])") - py.exec("res = jnp.linalg.inv(a1a2)") - a1a2.inv - } - } diff --git a/examples/src/main/scala/basic/LogisticRegression.scala b/examples/src/main/scala/basic/LogisticRegression.scala index badc2d16..b7cf5ac5 100644 --- a/examples/src/main/scala/basic/LogisticRegression.scala +++ b/examples/src/main/scala/basic/LogisticRegression.scala @@ -45,8 +45,8 @@ object LogisticRegression: case 0 => false } - val dataUnnormalized = Tensor2.fromArray(Axis[Sample], Axis[Feature], VType[Float])(featureData) - val dataLabels = Tensor1.fromArray(Axis[Sample], VType[Boolean])(labelData) + val dataUnnormalized = Tensor2(Axis[Sample], Axis[Feature]).fromArray(featureData) + val dataLabels = Tensor1(Axis[Sample]).fromArray(labelData) // TODO implement split val (trainingDataUnnormalized, valDataUnnormalized) = (dataUnnormalized, dataUnnormalized) @@ -88,7 +88,6 @@ object LogisticRegression: val trainLoss = loss(trainingData) val valLoss = loss(valData) val learningRate = 3e-1f - val xxx = summon[FloatTensorTree[BinaryLogisticRegression.Params]] val gd = GradientDescent(learningRate) val trainTrajectory = gd.iterate(initParams)(Autodiff.grad(trainLoss)) diff --git a/examples/src/main/scala/basic/Playground.scala b/examples/src/main/scala/basic/Playground.scala deleted file mode 100644 index 56e04481..00000000 --- a/examples/src/main/scala/basic/Playground.scala +++ /dev/null @@ -1,575 +0,0 @@ -package examples.basic - -import dimwit.* -import scala.util.NotGiven -import dimwit.random.Random -import dimwit.stats.Normal - -abstract class As[V, BaseType](using base: ExecutionType[BaseType]) extends ExecutionType[V]: - def dtype: DType = base.dtype - given ExecutionType[V] = this - -opaque type Y = Float - -trait A derives Label -trait B derives Label -trait C derives Label -trait D derives Label - -@main def playground(): Unit = - - trait Batch derives Label - trait Batch2 derives Label - trait Features derives Label - trait Samples derives Label - - val t = Tensor.zeros( - Shape( - Axis[Batch] -> 4, - Axis[Features] -> 8 - ), - VType[Y] - ) - - val t2 = Tensor.fromArray( - Shape( - Axis[Batch] -> 4, - Axis[Features] -> 8 - ), - VType[Float] - )(Array.fill(32)(1.0f)) - val t3 = t + t - val t4 = t2 + t2 - // val t5 = t + t2 // TODO this should not work - - println("TensorV2 Playground") - { - println("Normalization example") - val values = Array( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 - ).map(_.toFloat) - val X = Tensor.fromArray( - Shape( - Axis[Samples] -> 10, - Axis[Features] -> 2 - ), - VType[Float] - )(values) - val means = X.vmap(Axis[Features])(_.mean) - val stds = X.vmap(Axis[Features])(_.std) - val Xnorm = X.vmap(Axis[Samples]) { (x) => - (x - means) / stds - } - println(Xnorm) - println(Xnorm.shape) - println(Xnorm.device) - println(Xnorm.dtype) - } - { - println("DType and Device tests") - val t = Tensor.zeros( - Shape( - Axis[Batch] -> 1024, - Axis[Features] -> 512 - ), - VType[Float] - ) - println(t.shape) - println(t.dtype) - println(t.asType(VType[Int]).dtype) - println(t.device) - println(t.toDevice(Device.CPU).device) - } - { - val x = Tensor.zeros( - Shape( - Axis[Features] -> 2 - ), - VType[Float] - ) - val A = Tensor.zeros( - Shape( - Axis[Samples] -> 50, - Axis[Features] -> 2 - ), - VType[Float] - ) - // val y1 = x.dot(Axis[A])(A) - val y1 = A.dot(Axis[Features])(x) - println(y1.shape) - // A.dot(Axis["lala"])(x) - // A.dot(Axis[Samples])(x) - val y2 = x.dot(Axis[Features])(A) - println(y2.shape) - val y3 = x.outerProduct(A) - println(y3.shape) - } - { - println("Einops rearrange tests") - trait Batch derives Label - trait Frame derives Label - trait BatchFrame derives Label - trait Width derives Label - trait Height derives Label - trait Channel derives Label - trait Pixel derives Label - - val X = Tensor.zeros( - Shape( - Axis[Batch] -> 32, - Axis[Frame] -> 64, - Axis[Width] -> 256, - Axis[Height] -> 256, - Axis[Channel] -> 3 - ), - VType[Float] - ) - val d = X.rearrange( - ( - Axis[Batch |*| Frame], - Axis[Width |*| Height], - Axis[Channel] - ) - ) - println(d.shape) - val e = d.relabelAll((Axis[Frame], Axis[Pixel], Axis[Channel])) - println(e.shape) - } - { - println("Einops rearrange with trait-based labels") - trait Batch derives Label - trait Frame derives Label - trait Width derives Label - trait Height derives Label - trait Channel derives Label - val X = Tensor.zeros( - Shape( - Axis[Batch] -> 32, - Axis[Frame] -> 64, - Axis[Width] -> 256, - Axis[Height] -> 256, - Axis[Channel] -> 3 - ), - VType[Float] - ) - val d = X.rearrange( - ( - Axis[Batch |*| Frame], - Axis[Width |*| Height], - Axis[Channel] - ) - ) - println(d.shape) - } - { - - import scala.util.NotGiven - def f[L1: Label, L2: Label, L3: Label]( - x: Tensor[(L1, L2), Float], - y: Tensor[(L2, L3), Float] - ): Tensor[(L1, L3, L2), Float] = - x.vmap(Axis[L1]) { xi => - y.vmap(Axis[L3]) { yi => - xi + yi - } - } - val z = f( - Tensor.zeros( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ), - Tensor.zeros( - Shape( - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - ) - println(z.shape) - } - { - def f(t1: Tensor[(A, C), Float], t2: Tensor[Tuple1[C], Float]): Tensor[Tuple1[A], Float] = - t1.dot(Axis[C])(t2) - val t1 = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[C] -> 2 - ), - VType[Float] - ) - val t2 = Tensor.ones( - Shape( - Axis[C] -> 2 - ), - VType[Float] - ) - println(f(t1, t2)) - println("vmap 2") - import scala.util.NotGiven - val x1 = Tensor.ones( - Shape( - Axis[B] -> 1, - Axis[A] -> 2, - Axis[C] -> 2 - ), - VType[Float] - ) - val x2 = Tensor.ones( - Shape( - Axis[B] -> 4, - Axis[C] -> 2 - ), - VType[Float] - ) - println("vmap - println in vmap") - val res = x2.vmap(Axis[B]) { xi2 => - println(s"\t ${xi2}") - xi2 - } - println(res) - } - { - def f[L1: Label, L2: Label, L3: Label, V](x: Tensor[(L1, L2), Float], y: Tensor[(L2, L3), Float]) = - x.vmap(Axis[L1]) { xi => - y.vmap(Axis[L3]) { yi => - xi + yi - } - } - println( - f( - Tensor.zeros( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ), - Tensor.zeros( - Shape( - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - ).shape - ) - } - - { - println("Ravel") - val res = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - .ravel - println(res.shape) - } - { - println("swapaxes") - val res = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - .swap(Axis[A], Axis[C]) - println(res.shape) - } - { - println("appendAxis / prependAxis") - val res = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - .appendAxis(Axis[D]) - println(res.shape) - val res2 = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - .prependAxis(Axis[D]) - println(res2.shape) - } - { - println("squeeze") - val res = Tensor - .ones( - Shape( - Axis[A] -> 1, - Axis[B] -> 3, - Axis[C] -> 1 - ), - VType[Float] - ) - .squeeze(Axis[A]) - println(res.shape) - val res2 = res.squeeze(Axis[C]) - println(res2.shape) - } - { - println("Slice") - val res = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - .slice( - Axis[B] -> 2 - ) - println(res.shape) - val res2 = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - .slice( - Axis[B] -> (0 to 1) - ) - println(res2.shape) - val res3 = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4, - Axis[D] -> 5 - ), - VType[Float] - ) - .slice( - ( - Axis[B] -> 2, - Axis[C] -> 3 - ) - ) - println(res3.shape) - } - { - println("zipvmap tests") - trait Batch derives Label - trait Asset derives Label - trait Region derives Label - trait Sector derives Label - trait Risk derives Label - - val x = Tensor.ones( - Shape( - Axis[Batch] -> 6, - Axis[Asset] -> 3, - Axis[Region] -> 5 - ), - VType[Float] - ) - - val y = Tensor.ones( - Shape( - Axis[Region] -> 5, - Axis[Batch] -> 6, - Axis[Sector] -> 4 - ), - VType[Float] - ) - - val z = Tensor.ones( - Shape( - Axis[Sector] -> 4, - Axis[Risk] -> 5, - Axis[Batch] -> 6 - ), - VType[Float] - ) - - val res = zipvmap(Axis[Batch])(x, y) { case (xi, yi) => - xi.sum + yi.sum - } - println(res.shape) - - val res2 = zipvmap(Axis[Batch])(x, y, z) { case (xi, yi, zi) => - xi.sum + yi.sum + zi.sum - } - println(res2.shape) - } - { - import dimwit.tensor.* // Assuming imports - - trait Batch derives Label - trait Asset derives Label - trait Region derives Label - trait Sector derives Label - trait Risk derives Label - - val x = Tensor.ones(Shape(Axis[Batch] -> 6, Axis[Asset] -> 3, Axis[Region] -> 5), VType[Float]) - val y = Tensor.ones(Shape(Axis[Region] -> 5, Axis[Batch] -> 6, Axis[Sector] -> 4), VType[Float]) - val z = Tensor.ones(Shape(Axis[Sector] -> 4, Axis[Risk] -> 5, Axis[Batch] -> 6), VType[Float]) - - val res = zipvmap(Axis[Batch])((x, y, z)) { case (xi, yi, zi) => - xi.sum + yi.sum + zi.sum - } - println(res.shape) - } - { - println("TensorWhere tests") - val x = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - val y = Tensor.zeros( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - val condition = Tensor - .zeros( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - .asType(VType[Boolean]) - val res = where(condition, x, y) - println(res.shape) - } - { - println("Diag") - val x = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - val res = x.diagonal - println(res.shape) - } - { - println("Set") - val x = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - .set( - ( - Axis[A] -> 1, - Axis[B] -> 2 - ) - )(Tensor0(42)) - println(x) - val v = Tensor1.fromArray(Axis[B], VType[Float])( - Array(100, 101, 102).map(_.toFloat) - ) - val x2 = Tensor - .ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3 - ), - VType[Float] - ) - .set( - Axis[A] -> 1 - )(v) - println(x2) - } - { - // attention mechanism example - def softmax[L: Label](tensor: Tensor1[L, Float]): Tensor1[L, Float] = - val expTensor = tensor.exp - val sumExp = expTensor.sum - expTensor.vmap(Axis[L]) { _ / sumExp } - - trait Value derives Label - trait Key derives Label - trait Query derives Label - trait Context derives Label - - case class Attention( - wk: Tensor2[Value, Key, Float], - wq: Tensor2[Value, Query, Float], - wv: Tensor2[Value, Prime[Value], Float] - ): - private trait AttnWeights derives Label - - def apply(x: Tensor2[Context, Value, Float]): Tensor2[Context, Value, Float] = - val k = x.dot(Axis[Value])(wk) - val q = x.dot(Axis[Value])(wq) - val v = x.dot(Axis[Value])(wv) - val dk = Tensor0(Math.sqrt(k.shape(Axis[Key])).toFloat) - val attnWeightsPrime = q - .dot(Axis[Query ~ Key])(k) - .vmap(Axis[Context])(attnRow => softmax(attnRow).relabelTo(Axis[AttnWeights])) - val resPrime = attnWeightsPrime.dot(Axis[AttnWeights ~ Context])(v) - resPrime.relabel(Axis[Prime[Value]] -> Axis[Value]) - - trait Batch derives Label - - val x = Tensor.ones(Shape(Axis[Batch] -> 32, Axis[Context] -> 128, Axis[Value] -> 64), VType[Float]) - val attention = Attention( - Tensor.ones(Shape(Axis[Value] -> 64, Axis[Key] -> 64), VType[Float]), - Tensor.ones(Shape(Axis[Value] -> 64, Axis[Query] -> 64), VType[Float]), - Tensor.ones(Shape(Axis[Value] -> 64, Axis[Prime[Value]] -> 64), VType[Float]) - ) - val newX = x.vmap(Axis[Batch])(attention(_)) - println(newX.shape) - } - { - val t1 = Tensor.ones( - Shape( - Axis[A] -> 2, - Axis[B] -> 3, - Axis[C] -> 4 - ), - VType[Float] - ) - val t2 = t1.appendAxis(Axis[D]) - // val t3 = t1.appendAxis(Axis[A]) // should not compile - def f[T <: Tuple: Labels, V](t: Tensor[T, V]) = - t.appendAxis(Axis[D]) - def f2[T <: Tuple: Labels, V](t: Tensor[T, V]) = - t.appendAxis(Axis[A]) - val t3 = f(t1) - println(t3.shape) - val t4 = f2(t1) - println(t4.shape) - } - { - val x = Normal.standardNormal(Shape(Axis[A] -> 3, Axis[B] -> 4)) - println(x) - } diff --git a/examples/src/main/scala/complex/GPT2.scala b/examples/src/main/scala/complex/GPT2.scala index 06ac59b3..dbfaf3aa 100644 --- a/examples/src/main/scala/complex/GPT2.scala +++ b/examples/src/main/scala/complex/GPT2.scala @@ -116,8 +116,8 @@ case class GPT2(params: GPT2Params) extends (Tensor2[Batch, Context, Int] => Ten def causalMasking(attnScores: Tensor2[Context, Prime[Context], Float]): Tensor2[Context, Prime[Context], Float] = val ctxLength = attnScores.shape(Axis[Context]) - val causalMask = tril(Tensor.ones(Shape((Axis[Context] -> ctxLength, Axis[Prime[Context]] -> ctxLength)), VType[Boolean])) - where(causalMask, attnScores, Tensor.const(attnScores.shape, attnScores.vtype)(Float.NegativeInfinity)) + val causalMask = tril(Tensor(Shape((Axis[Context] -> ctxLength, Axis[Prime[Context]] -> ctxLength))).fill(true)) + where(causalMask, attnScores, Tensor.like(attnScores).fill(Float.NegativeInfinity)) val queries = x.dot(Axis[Embedding])(wq) +! wqBias val keys = x.dot(Axis[Embedding])(wk) +! wkBias @@ -209,10 +209,9 @@ case class Inference(gpt2: GPT2, tokenizer: Tokenizer): def loop(currentTokenIds: List[Int]): LazyList[String] = println(s"Current Token Ids: $currentTokenIds") val paddedTokenIds = currentTokenIds ++ List.fill(1024 - currentTokenIds.length)(0) - val inputTensor = Tensor.fromArray( - Shape((Axis[Batch] -> 1, Axis[Context] -> paddedTokenIds.length)), - VType[Int] - )( + val inputTensor = Tensor( + Shape((Axis[Batch] -> 1, Axis[Context] -> paddedTokenIds.length)) + ).fromArray( paddedTokenIds.toArray ) val predTokensTensor = gpt2(inputTensor).slice(Axis[Batch] -> 0) diff --git a/examples/src/main/scala/dataset/MNISTLoader.scala b/examples/src/main/scala/dataset/MNISTLoader.scala index 08b4f2c7..1acf50d4 100644 --- a/examples/src/main/scala/dataset/MNISTLoader.scala +++ b/examples/src/main/scala/dataset/MNISTLoader.scala @@ -37,7 +37,11 @@ object MNISTLoader: file.readFully(pixels) val shape = Shape(Axis[S] -> numImages, Axis[Height] -> rows, Axis[Width] -> cols) - Tensor.fromArray(shape)(pixels) + + // MNIST pixels are unsigned bytes + // So we read them as Byte and interpret as UInt8 when creating the Tensor + given ExecutionType[Byte] = ExecutionTypeFor[Byte](DType.UInt8) + Tensor(shape).fromArray(pixels) finally file.close() @@ -57,7 +61,7 @@ object MNISTLoader: file.readFully(labels) val shape = Shape(Axis[S] -> numLabels) - Tensor.fromArray(shape)(labels) + Tensor(shape).fromArray(labels) finally file.close() diff --git a/nn/src/main/scala/nn/Activation.scala b/nn/src/main/scala/nn/Activation.scala index ea95a6ea..878abd56 100644 --- a/nn/src/main/scala/nn/Activation.scala +++ b/nn/src/main/scala/nn/Activation.scala @@ -5,15 +5,11 @@ import dimwit.jax.Jax object ActivationFunctions: - // TODO rewrite relu, sigmoid to JAX - - def sigmoid[T <: Tuple: Labels](t: Tensor[T, Float]): Tensor[T, Float] = - val ones = Tensor.ones(t.shape, t.vtype) - ones / (ones + (-t).exp) + def sigmoid[T <: Tuple: Labels, V](t: Tensor[T, V]): Tensor[T, V] = + Tensor(Jax.jnn.sigmoid(t.jaxValue)) def relu[T <: Tuple: Labels, V](t: Tensor[T, V]): Tensor[T, V] = - val zeros = Tensor.zeros(t.shape, t.vtype) - maximum(t, zeros) + Tensor(Jax.jnn.relu(t.jaxValue)) def gelu[T <: Tuple: Labels, V](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnn.gelu(t.jaxValue)) diff --git a/nn/src/main/scala/nn/GradientOptimizer.scala b/nn/src/main/scala/nn/GradientOptimizer.scala index c6d95698..d0630dc6 100644 --- a/nn/src/main/scala/nn/GradientOptimizer.scala +++ b/nn/src/main/scala/nn/GradientOptimizer.scala @@ -68,7 +68,7 @@ case class Lion(learningRate: Tensor0[Float], weightDecay: Tensor0[Float] = Tens [T <: Tuple] => (n: Labels[T]) ?=> (t: Tensor[T, Float]) => - Tensor.zeros(t.shape, VType[Float]) + Tensor(t.shape).fill(0f) ) def update[Params: ToPyTree: FloatTensorTree](gradients: Grad[Params], params: Params, momentums: Params): (Params, Params) = diff --git a/nn/src/main/scala/nn/LinearLayer.scala b/nn/src/main/scala/nn/LinearLayer.scala index 0936d7f7..d457230e 100644 --- a/nn/src/main/scala/nn/LinearLayer.scala +++ b/nn/src/main/scala/nn/LinearLayer.scala @@ -23,7 +23,7 @@ object LinearLayer: ): Params[In, Out] = Params( weight = Normal.standardNormal(Shape(inputDim, outputDim)).sample(paramKey), - bias = Tensor.zeros(Shape(outputDim), VType[Float]) + bias = Tensor(Shape(outputDim)).fill(0.0f) ) case class LinearLayer[In: Label, Out: Label](params: LinearLayer.Params[In, Out]) extends Function[Tensor1[In, Float], Tensor1[Out, Float]]: