Extend eye method by two axes parameters.

[benikm91]

In my experiments, I ran into an issue: I had a shape2 mask and wanted to create an identity matrix from it, which was only possible with extra steps given the current eye implementation.

This PR adds more eye overloads to support more use cases. Specifically, it adds a use case for providing two axis extents to either create a square matrix with different labels (rather than L, Prime[L]). This also supports non-square matrix creation, mirroring what the jax eye method does. Additionally, I added an overload for Shape2, since it feels natural to create an eye matrix from a Shape2.

All eye fall back to Float32 VType, as the current eye implementation already does.

[marcelluethi]

I think the addition is useful. I always found the usage of eye to be somewhat awkward and confusing.

With these addition, I think it would be better to have eye as a method on the Tensor Factory method (like fill)

[marcelluethi]

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.

[marcelluethi]

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]

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]

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]

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.