WIP: Suggestion new tensor creation

[benikm91]

We discussed that the current Tensor creation API can be improved.

This PR suggests a new API similar to what we discussed.
Additionally, it provides more flexibility to create tensors with various dtypes.

---

New API
Change position of shape Tensor.

Old:
Tensor.fromArray(shape)(values)
New:
Tensor(shape).fromArray(array)

Replaced const, zeros, ones with fill:

// Old:
Tensor.ones(shape, VType[Float])
// New:
Tensor(shape).fill(1f)

// Old:
Tensor.zeros(shape, VType[Float])
// New:
Tensor(shape).fill(0f) // note 0 would result in an Int Tensor

// Old:
Tensor.const(shape, VType[Int])(42)
// New:
Tensor(shape).fill(42) // note 42f would result in a Float Tensor

// Old:
Tensor.zeros(shape, VType[Boolean])
// New:
Tensor(shape).fill(false)  // having false here instead of zeros is a big improvement (IMO)

// Similar to old API we have Tensor1, Tensor2 and Tensor3 syntax sugars for fromArray
Tensor1(Axis[A]).fromArray(values)  // length of A is taken from the values: Array[V]
Tensor2(Axis[A], Axis[B]).fromArray(values)  // length of A is taken from the values: Array[Array[V]]

// Tensor0 has an apply syntax sugar (as before) instead of Tensor0.fill(...)
Tensor0(0f)   // Tensor[EmptyTuple, Float]

Fix Execution Type to allow dtype definitions in fill

We might have to discuss the API for chaning dtype (see below), as
Tensor(shape).fromArray(pixels, dtype=DType.UInt8)
would seem more natural. However, I suggest moving this discussion to a later stage and focusing on the "default-dtype" API for now. The motivation for a type class is to change the behavior globally, e.g., switching from float32 to float64 in a simulation...

However, this PR did not change the ExecutionType API, just cleaned up its implementation so it works with the new fill operation (zeros, ones, const did not support changing dtypes). One minor change was changing Array[Byte] to be mapped to int8 by default not uint8, to be in line what the Scala type means.

// Current dtype changing example:
val t = Tensor1(Axis[A] -> 2).fill(3.14)  // Type is Tensor[?, Float]
t.dtype shouldBe DType.Float64
given ExecutionType[Double] = ExecutionTypeFor[Double](DType.Float32) // Overwrite dtype globally
val t = Tensor1(Axis[A] -> 2).fill(3.14)
t.dtype shouldBe DType.Float32

---

Some conceptual things:

• ExecutionType is a temporary object on the Scala sight. It defines only at the new tensor creation what the precision of the tensor will be. Afterwards it is removed, i.e., the tensor does not store the ExecutionType. A tensorlooks the dtype up in JAX if requested. This is due to operations changing the dtype in JAX (float32 * float64 -> float64), which we otherwise would have to track.
• V vs dtype: The Tensor value type V is either Float, Int, Boolean, or an opaque type. We do ignore precision on purpose here, as it is not conceptual, only about execution. Same as we do not track the device. However, this can be confusing. For example, if I cast the MNIST images to int8 (the default for Byte), it results in weird runtime behavior.

---

Code changes to check
As most files have been touched I highlight the main changes here:

1. Playground and TensorAPI and AutoDiffAPI are deleted, as we know have tests (didn't want to translate them to new API).
2. Tensor.scala contains new API definitions
3. DType now knows based on itself how to write an Array[?] to an Array of Bytes. This logic was before in Tensor.scala, and less flexible. It is required in Tensor.scala.
4. ArrayWriter and WriterEvidence provide a tool to translate a scalar A or an array Array[A] into a Tensor[?, V]. Note that A and V are different types, as we can cast a Double into a Tensor[?, Float]. The underlying precision in JAX depends on the type A, so you can specify, if Doubles should be represented in Float64 or Float32. Or if Bytes should be represented in Int8 or UInt8
5. New TestSuite TensorCreationSuite showing how the creation and overwriting the precision works.

[marcelluethi]

I like the new interface a lot. It makes things both cleaner and more terse.

One thing I was wondering tough:

    val t1 = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1f, 2f, 3f))
    val t2 = Tensor(Shape(Axis[A] -> 3)).fill(1.0)
    val t3 = Tensor(Shape(Axis[A] -> 3)).fill(true)
    // why
    val t4 = Tensor.randn(Shape(Axis[A] -> 3))(Random.Key(42))
    // and not
    val t5 = Tensor(Shape(Axis[A] -> 3)).randn(Random.Key(42))

The last line could be easily implemented by just moving the randn function to the factory. Is there a reason why you did not choose to do that?

[benikm91]

I will add this.

randn is something I want to change (to support any distribution). However, this could also be part of future changes. So thats why I overlooked it.

So I will update randn.

[marcelluethi]

Great - if failing tests are addressed and randn is supported, it seems ready to merge.

Regarding supporting all distributions. I think it is good to postpone that, as this doesn't seem to be easy. Currently distributions need a shape to be constructed. Once the shape of a distribution is fixed, sampling is easy - so we wouldn't even need convenience functions. Maybe this is the reasons that most libraries just allow randn (standard normal) and rand (uniform 0, 1)

[benikm91]

We currently use randn nowhere (except for Tensor0 in RandomSuite).

I would suggest removing "randn" for now and keeping the Normal.standardNormal(shape).sample(key) syntax.

// Current
weight = Normal.standardNormal(Shape(inputDim)).sample(paramKey),
// randn
weight = Tensor(Shape(inputDim)).randn(paramKey),