Idea: ∂ as label modifier

[benikm91]

(Read critically. This is an initial idea and could be flawed)

For grad we have a Grad[] type to make the gradient type different from the Params type, e.g.

val params: Params = ???
val grad: Grad[Params] = Autodiff.grad(f)(params)

Such a distinction is missing for jacobian and hessian, leading to them having the same type for different concepts:

val f: Tensor1[A] => Tensor0
val g: Tensor1[A] => Tensor1[A]
// hessian(f) and jacobian(g) have the same type Tensor1[A] => Tensor2[A, Prime[A]] despite being semantically different

For the hessian it is also unclear what A and Prime[A] mean.

This thinking led me to the following insight. Grad is not something that should wrap the tensor but being a partial derivative ∂ is a concept of a tensor dimension. This would mean we could express things like this:

f: Tensor1[A] => Tensor0
df: Tensor1[A] => Tensor1[∂A]
ddf: Tensor1[A] => Tensor2[∂A, ∂A]

f: Tensor1[A] => Tensor1[B]
df: Tensor1[A] => Tensor2[B, ∂A]
ddf: Tensor1[A] => Tensor3[B, ∂A, ∂A]

f: (a: Tensor1[A], b: Tensor1[B]) => Tensor1[C]
df: (a: Tensor1[A], b: Tensor1[B]) => (
    a: Tensor2[C, ∂A],
    b: Tensor2[C, ∂B]
)
ddf: (a: Tensor1[A], b: Tensor1[B]) => (
    a: (
        a: Tensor3[C, ∂A, ∂A],
        b: Tensor3[C, ∂A, ∂B]
    ),
    b: (
        a: Tensor3[C, ∂B, ∂A],
        b: Tensor3[C, ∂B, ∂B]
    )
)

This would clarify the meaning of the dimensions and distinct hessian(f) and jacobian(g) from before.

hessian(f) // Tensor1[A] => Tensor2[∂A, ∂A],
jacobian(g) // Tensor1[A] => Tensor2[A, ∂A],

We could also express grad with this, however it comes at the cost of requiring named tuples instead of case classes as case classes have fixed labels which we can't overwrite with ∂. Grad[Params] can also be kept (change opaque type to new ∂ mechanism internally).

type Params = (
  w: Tensor2[A, B, Float32],
  b: Tensor1[A, Float32],
)
val f: Params => Tensor0
val df = Autodiff.grad(f) // Params => (w: Tensor2[∂A, ∂B, Float32], b: Tensor1[∂A, Float32])