sbrunk/storch

Tensor.grad should return an option

Closed this issue · 0 comments

sbrunk commented

The tensor returned by grad is only defined after a backward() call. Right now, we're returning an undefined tensor if we don't have gradients:

def grad: Tensor[D | Undefined] = Tensor(native.grad())
import torch.*

val t = torch.rand(Seq(1), requiresGrad=true)
// t: Tensor[Float32] = tensor dtype=float32, shape=[1], device=CPU 
// [0,0102]
t.grad
// res0: Tensor[Float32 | Undefined] = undefined tensor
// t.grad + 1 // boom!

t.backward()
t.grad
// res2: Tensor[Float32 | Undefined] = tensor dtype=float32, shape=[1], device=CPU 
// [1,0000]

That's somewhat unsafe since we can call any operation on the gradient tensor and get a runtime error if it's undefined. But it is also awkward to handle if defined, because we need to explicitly cast it from Tensor[Float32 | Undefined] to its actual type (Tensor is not covariant).

So instead, we could return an Option[Tensor[D], containing the value or None depending on whether it is defined or not.