Enable scalars if compiled with WITH_SCALAR environment variable. (#4806)

* Enable scalars if compiled with WITH_SCALAR environment variable.

We are pretty close to enabling scalars (0-dimensional arrays); this allows turning them on
for development purposes and to be able to write code that works both with and without scalars enabled.

WITH_SCALARS is currently broken with distributions, but should work for test_torch, test_autograd, test_nn.

* Fix unsqueeze.

* Fix wrap dim, wrapping with Scalar.

5 files changed, 37 insertions, 35 deletions

diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
index 2b5790eaf0..56ddb55f12 100644
--- a/tools/autograd/derivatives.yaml
+++ b/tools/autograd/derivatives.yaml
@@ -552,7 +552,7 @@
   self: unsqueeze_to(grad, self.sizes());
 
 - name: squeeze(Tensor self, int64_t dim)
-  self: maybe_unsqueeze(grad, dim, self.size(dim) == 1 && self.sizes().size() != 1)
+  self: unsqueeze_to(grad, dim, self.sizes())
 
 - name: std(Tensor self, bool unbiased)
   self: var_backward(grad / (result * 2), self, unbiased)
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
index 672d6c6e0a..b20166751d 100644
--- a/tools/autograd/gen_variable_type.py
+++ b/tools/autograd/gen_variable_type.py
@@ -23,6 +23,7 @@
 #     differentiable subcomponents.
 #
 from __future__ import print_function
+import os
 import sys
 from .utils import CodeTemplate, nested_dict, write
 from .gen_autograd import VIEW_FUNCTIONS, template_path
@@ -458,7 +459,7 @@ def emit_body(declaration):
     body.append(declare_returned_variables())
     body.append(emit_call(env))
     if requires_derivative:
-        if inplace and is_view:
+        if inplace and is_view and not os.environ.get('WITH_SCALARS'):
             body.append('ensure_no_aten_scalars(self);')
         # set_flags has to appear after version_counter, because rebase_history
         # requires that the counter is incremented before it is called
diff --git a/tools/autograd/templates/Functions.cpp b/tools/autograd/templates/Functions.cpp
index 2e56f80ddd..e6732ee0b5 100644
--- a/tools/autograd/templates/Functions.cpp
+++ b/tools/autograd/templates/Functions.cpp
@@ -65,7 +65,11 @@ Tensor norm_backward(const Tensor & grad, const Tensor & self, const Scalar & p_
 }
 
 Tensor norm_backward(Tensor grad, const Tensor & self, const Scalar & p_, Tensor norm, int64_t dim, bool keepdim) {
+#ifdef WITH_SCALARS
+  if (!keepdim) {
+#else
   if (!keepdim && self.dim() > 1) {
+#endif
     grad = grad.unsqueeze(dim);
     norm = norm.unsqueeze(dim);
   }
@@ -98,7 +102,11 @@ Tensor permute_backwards(const Tensor & grad, IntList fwd_dims) {
 }
 
 Tensor sum_backward(const Tensor & grad, IntList sizes, int64_t dim, bool keepdim) {
-  if (!keepdim && sizes.size() > 1) {
+#ifdef WITH_SCALARS
+  if (!keepdim) {
+#else
+   if (!keepdim && sizes.size() > 1) {
+#endif
     return grad.unsqueeze(dim).expand(sizes);
   } else {
     return grad.expand(sizes);
@@ -302,6 +310,7 @@ Tensor cumsum_backward(const Tensor & x, int64_t dim) {
 Tensor unsqueeze_to(const Tensor & self, IntList sizes) {
   auto result = self;
 
+#ifndef WITH_SCALARS
   // Let's say the input had size (1, 1). input.squeeze(), with scalars
   // disabled, produces a result of size (1,). This needs some
   // special handling because for all other cases we unsqueeze every
@@ -310,6 +319,7 @@ Tensor unsqueeze_to(const Tensor & self, IntList sizes) {
   if (self.sizes().equals({1})) {
     return result.view(sizes);
   }
+#endif
 
   int64_t nDims = sizes.size();
   for (int64_t dim = 0; dim < nDims; dim++) {
@@ -320,8 +330,13 @@ Tensor unsqueeze_to(const Tensor & self, IntList sizes) {
   return result;
 }
 
-Tensor maybe_unsqueeze(const Tensor & self, int64_t dim, bool unsqueeze) {
-  if (unsqueeze) {
+Tensor unsqueeze_to(const Tensor & self, int64_t dim, IntList sizes) {
+  dim = at::maybe_wrap_dim(dim, sizes.size());
+#ifdef WITH_SCALARS
+  if (sizes[dim] == 1) {
+#else
+  if (sizes[dim] == 1 && sizes.size() != 1) {
+#endif
     return self.unsqueeze(dim);
   }
   return self;
@@ -387,14 +402,22 @@ Tensor renorm_backward(const Tensor & grad, const Tensor & self, Scalar p, int64
 }
 
 Tensor select_backward_scalar(Tensor grad, const Tensor & input, const Tensor & value) {
-  auto grad_data = static_cast<Variable&>(grad).data();
   auto grad_input = zeros_like(input);
+#ifdef WITH_SCALARS
+  grad_input.masked_fill_(input == value, grad);
+#else
+  auto grad_data = static_cast<Variable&>(grad).data();
   grad_input.masked_fill_(input == value, Scalar(grad_data[0]));
+#endif
   return grad_input;
 }
 
 Tensor select_backward(Tensor grad, int64_t dim, Tensor indices, IntList sizes, bool keepdim) {
+#ifdef WITH_SCALARS
+  if (!keepdim) {
+#else
   if (!keepdim && sizes.size() > 1) {
+#endif
     grad = grad.unsqueeze(dim);
     indices = indices.unsqueeze(dim);
   }
@@ -406,13 +429,16 @@ Tensor trace_backward(const Tensor & grad, IntList sizes) {
     throw std::runtime_error("expected matrix input");
   }
 
-  // TODO: simplify once index_fill_(Tensor) is implemented on Variable
-  auto grad_data = static_cast<const Variable&>(grad).data();
   auto& long_type = grad.type().toScalarType(at::kLong);
 
   auto grad_input = grad.type().zeros(sizes[0] * sizes[1]);
   auto indices = long_type.arange(0, grad_input.numel(), sizes[1] + 1);
+#ifdef WITH_SCALARS
+  grad_input.index_fill_(0, indices, grad);
+#else
+  auto grad_data = static_cast<const Variable&>(grad).data();
   grad_input.index_fill_(0, indices, Scalar(grad_data[0]));
+#endif
   return grad_input.view(sizes);
 }
 
diff --git a/tools/autograd/templates/VariableType.cpp b/tools/autograd/templates/VariableType.cpp
index 7b9e07df56..5221d722b9 100644
--- a/tools/autograd/templates/VariableType.cpp
+++ b/tools/autograd/templates/VariableType.cpp
@@ -239,11 +239,13 @@ static Tensor as_view(const Tensor & base, Tensor tensor) {
   return make_variable_view(std::move(base_var), std::move(tensor));
 }
 
+#ifndef WITH_SCALARS
 static void ensure_no_aten_scalars(Tensor & data) {
   if (data.defined() && data.dim() == 0) {
     data.as_strided_({1}, {1});
   }
 }
+#endif
 
 template<typename T>
 static bool computes_grad_tmpl(T tensors) {
diff --git a/tools/autograd/templates/python_variable_methods.cpp b/tools/autograd/templates/python_variable_methods.cpp
index f519d1002b..2110cb260d 100644
--- a/tools/autograd/templates/python_variable_methods.cpp
+++ b/tools/autograd/templates/python_variable_methods.cpp
@@ -523,32 +523,6 @@ static PyObject * THPVariable_type(PyObject* self, PyObject* args, PyObject* kwa
   END_HANDLE_TH_ERRORS
 }
 
-// FixMe: remove when scalars fully supported
-inline PyObject* _wrap_scalar(at::Tensor tensor) {
-  if (!tensor.sizes().equals({1})) {
-    throw std::runtime_error("tried to wrap scalar of non-scalar size");
-  }
-  auto v = Variable(std::move(tensor));
-  v.data().squeeze_();
-  return THPVariable_Wrap(v, true);
-}
-
-static PyObject * THPVariable__scalar_sum(PyObject* self, PyObject* args, PyObject* kwargs)
-{
-  HANDLE_TH_ERRORS
-  static PythonArgParser parser({
-    "sum()",
-  });
-  auto& self_ = reinterpret_cast<THPVariable*>(self)->cdata;
-  PyObject* parsed_args[3];
-  auto r = parser.parse(args, kwargs, parsed_args);
-  if (r.idx == 0) {
-    return _wrap_scalar(dispatch_sum(self_));
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
 // generated methods start here
 
 ${py_methods}
@@ -606,7 +580,6 @@ PyMethodDef variable_methods[] = {
   {"stride", (PyCFunction)THPVariable_stride, METH_VARARGS | METH_KEYWORDS, NULL},
   {"tolist", (PyCFunction)THPVariable_tolist, METH_NOARGS, NULL},
   {"type", (PyCFunction)THPVariable_type, METH_VARARGS | METH_KEYWORDS, NULL},
-  {"_scalar_sum", (PyCFunction)THPVariable__scalar_sum,  METH_VARARGS | METH_KEYWORDS, NULL},
   ${py_method_defs}
   {NULL}
 };