error instead of crashing on attempt to subclass typed tensors (#20283)

nairbv · facebook-github-bot · commit c406bf20a00b · 2019-05-09T07:10:38.000-07:00
Summary: #20052 typed tensors (e.g. torch.FloatTensor) can't be subclassed. Was causing crashes and other errors. Pull Request resolved: #20283 Differential Revision: D15278138 Pulled By: nairbv fbshipit-source-id: 8493eac4d34dfb76b054362bf0acec02146cd0e2
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -11336,6 +11336,19 @@ def test_c10_layer_norm(self):
                 weight), torch.tensor(bias), 1, epsilon, True)
         torch.testing.assert_allclose(expected_norm, actual_norm)
 
+    def test_subclass_tensors(self):
+        # raise an error when trying to subclass FloatTensor
+        with self.assertRaisesRegex(TypeError, "type 'torch.FloatTensor' is not an acceptable base type"):
+            class Foo1(torch.FloatTensor):
+                pass
+
+        # but allow subclassing Tensor:
+        class Foo2(torch.Tensor):
+            def foo(self):
+                return 5
+        f = Foo2()
+        self.assertEqual(f.foo(), 5)
+
 # Functions to test negative dimension wrapping
 METHOD = 1
 INPLACE_METHOD = 2
diff --git a/torch/csrc/tensor/python_tensor.cpp b/torch/csrc/tensor/python_tensor.cpp
@@ -159,7 +159,9 @@ static void py_initialize_tensor_type(PyTypeObject& type, const char* name, PyOb
   ((PyObject*)&type)->ob_refcnt = 1;
   ((PyObject*)&type)->ob_type = &metaclass;
   type.tp_basicsize = sizeof(PyTensorType);
-  type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
+  // Subclassing from torch.<ScalarType>Tensor isn't supported.
+  // (Py_TPFLAGS_BASETYPE omitted). Subclassing torch.Tensor still allowed.
+  type.tp_flags = Py_TPFLAGS_DEFAULT;
   type.tp_name = name;
   type.tp_new = Tensor_new;
   if (PyType_Ready(&type) < 0) {
