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.