Skip to content

Extend eye method by two axes parameters. #160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
benikm91 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Extend eye method by two axes parameters. #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions core/src/main/scala/dimwit/tensor/Tensor.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,9 +330,15 @@ object Tensor2:
def apply[L1: Label, L2: Label](axisExtent1: AxisExtent[L1], axisExtent2: AxisExtent[L2]): Tensor.DefaultsFactory[Tuple2[L1, L2]] = Tensor.DefaultsFactory(Shape(axisExtent1, axisExtent2))
def apply[L1: Label, L2: Label, V](axisExtent1: AxisExtent[L1], axisExtent2: AxisExtent[L2], vtype: VType[V]): Tensor.TypedFactory[Tuple2[L1, L2], V] = Tensor.TypedFactory(Shape(axisExtent1, axisExtent2), vtype)

private def eyeImpl[L: Label, V](dim: AxisExtent[L], vtype: VType[V]): Tensor2[L, Prime[L], V] = Tensor(Jax.jnp.eye(dim.size, dtype = vtype.dtype.jaxType))
def eye[L: Label](dim: AxisExtent[L]): Tensor2[L, Prime[L], Float32] = eyeImpl(dim, VType[Float32])
def eye[L: Label, V](dim: AxisExtent[L], vtype: VType[V]): Tensor2[L, Prime[L], V] = eyeImpl(dim, vtype)
private def eyeImpl[T <: Tuple: Labels, V](n: Int, m: Int, vtype: VType[V] = VType[Float32]): Tensor[T, V] = Tensor(Jax.jnp.eye(n, m, dtype = vtype.dtype.jaxType))

def eye[L: Label](dim: AxisExtent[L]): Tensor2[L, Prime[L], Float32] = eye(dim, Axis[Prime[L]] -> dim.size)

@marcelluethi marcelluethi Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now that we have so many variants of eye, we should comment them (short common including purpose) making it easier for the user to choose the correct variant.

def eye[L: Label, V](dim: AxisExtent[L], vtype: VType[V]): Tensor2[L, Prime[L], V] = eye(dim, Axis[Prime[L]] -> dim.size, vtype)

@marcelluethi marcelluethi Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it make sense to add eye as a factory method, similar to the others. The user would then write

Tensor2(Axis[A]->3, Axis[B]->2).eye
instead of
Tensor2.eye(Axis[A]-3, Axis[B]->2

@benikm91 benikm91 Sep 2, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I looked into this a little bit.

First, implementation is a little bit more complex than one might expect, as Tensor2(Axis[A]->3, Axis[B]->2) currently returns a Tensor.DefaultsFactory; we could implement a Tensor2.DefaultsFactory that wraps all methods in Tensor.DefaultsFactory (or extends?) and additionally implements an .eye method. Putting eye in Tensor.DefaultsFactory as it would allow Tensor1(Axis[A]->3).eye.

This implementation detail to the side and think about the API only. Let's lay out what we currently have:

Current API:

Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) // Float Array -> Float32
Tensor2(Axis[A] -> 3, Axis[B] -> 2).fromArray(1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f) // Float Array -> Float32
Tensor2(Axis[A] -> 3, Axis[B] -> 2).fill(32) // Int Value -> Int32

Eye:

 // Current; 
Tensor2.eye(Axis[A] -> 3, Axis[B] -> 3) // Default -> Float32
Tensor2.eye(Axis[A] -> 3, Axis[B] -> 3, VType[Int32]) // Int32
// New
Tensor2(Axis[A] -> 3, Axis[B] -> 3).eye // Default -> Float32
Tensor2(Axis[A] -> 3, Axis[B] -> 3).eye(VType[Int32]) // Int32

One advantage that the current implementation has in eye is that we can import Tensor2.eye to use eye as a method, like in JAX.

Another (currently missing) function similar to eye would be Tensor1.range, a special constructor method for a specific-rank tensor.

Tensor1.range(axis, start=0, stop=10, step=1) // Currrent, Int32
Tensor1(axis).range(start=0, stop=10, step=1) // New, Int32

Here, the New syntax is the obvious right choice for me. So I guess for eye it must be the same. The VType could here be set by the types of the arguments

Tensor1.range(axis, start=0f, stop=10f, step=1f) // Currrent, Float32
Tensor1(axis).range(start=0f, stop=10f, step=1f) // New, Float32

=> My final conclusion would be that I like your suggestion, and I would pick the Tensor().eye syntax.

If we go with a Tensor2.DefaultsFactory, I would additionally provide the following method. The New method would check if the 2-dimensional array input matches the expected shape. This lets the user explicitly state the expected shape. "Exists" would still be available where the shape is defined by the input parameter.

Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) // Exists
Tensor2(Axis[A] -> 2, Axis[B] -> 2).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) // New

@marcelluethi marcelluethi Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for looking into this. I think we agree then on a unified principle to construct new tensors. First the Axis/Shapes are fixed and then we have the methods like eye, range, fromArray working on the tensors with a fixed size.

Regarding this example:

Tensor2(Axis[A] -> 2, Axis[B] -> 2).fromArray(Array(Array(1.0f, 2.0f), Array(3.0f, 4.0f))) // New

Would this be checked at compile time? If not, what is the advantage to the already existing method

Tensor2(Axis[A] -> 2, Axis[B] -> 2).fromArray(Array(1.0f, 2.0f, 3.0f, 4.0f)) // Exists
``

@benikm91 benikm91 Sep 3, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just in case we want to validate the correct shape and we already have a multidimensional array in Scala.

val mnistDataset: Array[Array[Float]] = ???
val pixelExtent = Axis[Pixel] -> 8*8
Tensor2(Axis[Image] -> 50_000, pixelExtent).fromArray(mnistDataset) // allow this
Tensor2(Axis[Image] -> 50_000, pixelExtent).fromArray(mnistDataset.flatten) //  ohterwise.

def eye[L1: Label, L2: Label](n: AxisExtent[L1], m: AxisExtent[L2]): Tensor2[L1, L2, Float32] = eyeImpl[(L1, L2), Float32](n.size, m.size)
def eye[L1: Label, L2: Label, V](n: AxisExtent[L1], m: AxisExtent[L2], vtype: VType[V]): Tensor2[L1, L2, V] = eyeImpl[(L1, L2), V](n.size, m.size, vtype)
def eye[L1: Label, L2: Label](shape: Shape2[L1, L2]): Tensor2[L1, L2, Float32] = eye(shape.extent(Axis[L1]), shape.extent(Axis[L2]))
def eye[L1: Label, L2: Label, V](shape: Shape2[L1, L2], vtype: VType[V]): Tensor2[L1, L2, V] = eye(shape.extent(Axis[L1]), shape.extent(Axis[L2]), vtype)

def diag[L: Label, V](diag: Tensor1[L, V]): Tensor2[L, Prime[L], V] = Tensor(Jax.jnp.diag(diag.jaxValue))

/** Companion object for Tensors of rank 3.
Expand Down
42 changes: 42 additions & 0 deletions core/src/test/scala/dimwit/tensor/TensorCreationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,45 @@ class TensorCreationSuite extends DimwitTest:
idx(Axis[A]).toFloat
}
result shouldEqual Tensor1(Axis[A]).fromArray(Array(0.0f, 1.0f, 2.0f, 3.0f))

describe("eye"):

it("square: from a single extent"):
val result = Tensor2.eye(Axis[A] -> 3)
result.shape shouldEqual Shape2(Axis[A] -> 3, Axis[Prime[A]] -> 3)
result shouldEqual Tensor2(Axis[A], Axis[B]).fromArray(
Array(Array(1.0f, 0.0f, 0.0f), Array(0.0f, 1.0f, 0.0f), Array(0.0f, 0.0f, 1.0f))
)

it("square: from two extents and from a shape"):
val expected = Tensor2(Axis[A], Axis[B]).fromArray(
Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f))
)
Tensor2.eye(Axis[A] -> 2, Axis[B] -> 2) shouldEqual expected
Tensor2.eye(Shape2(Axis[A] -> 2, Axis[B] -> 2)) shouldEqual expected

it("wide: more columns than rows, zero padded"):
val expected = Tensor2(Axis[A], Axis[B]).fromArray(
Array(Array(1.0f, 0.0f, 0.0f), Array(0.0f, 1.0f, 0.0f))
)
val fromExtents = Tensor2.eye(Axis[A] -> 2, Axis[B] -> 3)
fromExtents.shape shouldEqual Shape2(Axis[A] -> 2, Axis[B] -> 3)
fromExtents shouldEqual expected
Tensor2.eye(Shape2(Axis[A] -> 2, Axis[B] -> 3)) shouldEqual expected

it("tall: more rows than columns, truncating"):
val expected = Tensor2(Axis[A], Axis[B]).fromArray(
Array(Array(1.0f, 0.0f), Array(0.0f, 1.0f), Array(0.0f, 0.0f))
)
val fromExtents = Tensor2.eye(Axis[A] -> 3, Axis[B] -> 2)
fromExtents.shape shouldEqual Shape2(Axis[A] -> 3, Axis[B] -> 2)
fromExtents shouldEqual expected
Tensor2.eye(Shape2(Axis[A] -> 3, Axis[B] -> 2)) shouldEqual expected

it("defaults to Float32 and accepts an explicit vtype in all variants"):
Tensor2.eye(Axis[A] -> 2).dtype shouldBe DType.Float32
Tensor2.eye(Axis[A] -> 2, Axis[B] -> 3).dtype shouldBe DType.Float32
Tensor2.eye(Shape2(Axis[A] -> 2, Axis[B] -> 3)).dtype shouldBe DType.Float32
Tensor2.eye(Axis[A] -> 2, VType[Int32]).dtype shouldBe DType.Int32
Tensor2.eye(Axis[A] -> 2, Axis[B] -> 3, VType[Int32]).dtype shouldBe DType.Int32
Tensor2.eye(Shape2(Axis[A] -> 2, Axis[B] -> 3), VType[Int32]).dtype shouldBe DType.Int32
Loading