Add variance() and variance(squeezingAxes:).

[dan-zheng]

• Add variance() and variance(squeezingAxes:).
• Add derivatives for squeezingAxes: versions of sum and mean.
• Test variance.
• Test variance derivatives.

Todos:

• Add "bias correction" flag to variance ops.

[dan-zheng]

Our variance ops are implemented using mean and squared - no custom derivatives for variance.

However, the derivatives don't seem to match Python:

import torch
import torch.nn.functional as F

x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float)
x.requires_grad = True
y = x.var(unbiased=False)
y.backward()

print(x)
# tensor([[1., 2.],
#         [3., 4.]], requires_grad=True)

print(y)
# tensor(1.2500, grad_fn=<VarBackward0>)

print(x.grad)
# tensor([[-0.7500, -0.2500],
#         [ 0.2500,  0.7500]])
# We are not consistent with this.

Any ideas on fixing on this? Is a custom variance derivative necessary?
(I would find it strange if a custom derivative is necessary for correctness.)

Seeking help: cc @jekbradbury @jph00 @sgugger

[dan-zheng]

PyTorch has a custom var_backward implementation:

Tensor var_backward(const Tensor & grad, const Tensor & self, bool unbiased) {
  return (2.0 / (self.numel() - unbiased)) * grad * (self - self.mean());
}

Tensor var_backward(Tensor grad, const Tensor & self, IntArrayRef dim, bool unbiased, bool keepdim) {
  if (self.dim() == 0) {
    return var_backward(grad, self, unbiased);
  }
  if (!keepdim && self.dim() > 1) {
    grad = unsqueeze_multiple(grad, dim, self.sizes().size());
  }
  return (2.0 / (_safe_size(self.sizes(), dim) - unbiased)) * grad * (self - self.mean(dim, true));
}

I'll try to port it to Swift.

I'm still concerned about derivative correctness - a custom derivative shouldn't be necessary for correctness, especially because variance isn't super complicated (it just calls mean and squared).

[jph00]

Sorry it's not something that @sgugger or I have looked at before so we don't have any particular insights here.

[rxwei]

Why don't we hand-calculate and verify a simple case? It'd be highly preferable to not make it a primitive differentiable function.

[sgugger]

PyTorch is correct and doing

let xGrad = gradient(at: x) {x in return (x - x.mean()).squared().mean()}

or

let xGrad = gradient(at: x) {x in return {x in return x.variance(alongAxes: 0,1)}}

gives me the same result as PyTorch.

[jekbradbury]

The PyTorch gradient is correct. This Swift version also gives the right results:

@differentiable
func variance(_ x: Tensor<Float>) -> Tensor<Float> {
    let mean = x.mean()
    let squaredDiff = (x - mean).squared()
    return squaredDiff.mean()
}

[dan-zheng]

I wonder if reduction ops still fail on TPU. The #if !TPU condition was added months ago (b/111815968).

cc @pschuh @bgogul: I wonder if you have any clue if this test still fails on TPU?
Or if toHost(shape: ...) is necessary everywhere?

[dan-zheng]

Also, the arguments to toHost(shape:) don't seem quite right (it shouldn't always be []).

[dan-zheng]

Update: I'll remove the toHost calls to force us to re-evaluate this bug.

[dan-zheng]

@swift-ci Please test tensorflow