Added .bfloat16() (pytorch#22852)

Summary:
Add conversion method for bfloat16
Pull Request resolved: pytorch#22852

Differential Revision: D16256760

Pulled By: izdeby

fbshipit-source-id: 01d75495f9df513a0cdf78791c3eb013ab92bd95

Author: izdeby

docs/source/tensors.rst

@@ -149,7 +149,7 @@ view of a storage and defines numeric operations on it.
    .. autoattribute:: device
    .. autoattribute:: grad
    .. autoattribute:: ndim
-   .. autoattribute:: T 
+   .. autoattribute:: T
 
    .. automethod:: abs
    .. automethod:: abs_
@@ -185,6 +185,7 @@ view of a storage and defines numeric operations on it.
    .. automethod:: baddbmm_
    .. automethod:: bernoulli
    .. automethod:: bernoulli_
+   .. automethod:: bfloat16
    .. automethod:: bincount
    .. automethod:: bitwise_not
    .. automethod:: bitwise_not_

test/test_torch.py

@@ -158,6 +158,7 @@ def test_type_conversion_via_dtype_name(self):
         self.assertEqual(x.float().dtype, torch.float32)
         self.assertEqual(x.half().dtype, torch.float16)
         self.assertEqual(x.int().dtype, torch.int32)
+        self.assertEqual(x.bfloat16().dtype, torch.bfloat16)
 
     def test_doc(self):
         checked_types = (types.MethodType, types.FunctionType,

tools/autograd/templates/python_variable_methods.cpp

@@ -415,6 +415,10 @@ static PyObject * THPVariable_bool(PyObject* self, PyObject* args) {
   return THPVariable_to_type(self, ScalarType::Bool);
 }
 
+static PyObject * THPVariable_bfloat16(PyObject* self, PyObject* args) {
+  return THPVariable_to_type(self, ScalarType::BFloat16);
+}
+
 static PyObject * THPVariable_element_size(PyObject* self, PyObject* args)
 {
   HANDLE_TH_ERRORS
@@ -731,6 +735,7 @@ PyMethodDef variable_methods[] = {
   {"__matmul__", (PyCFunction)THPVariable_matmul, METH_VARARGS | METH_KEYWORDS, NULL},
   {"_is_view", (PyCFunction)THPVariable__is_view, METH_NOARGS, NULL},
   {"apply_", (PyCFunction)THPVariable_apply_, METH_O, NULL},
+  {"bfloat16", (PyCFunction)THPVariable_bfloat16, METH_NOARGS, NULL},
   {"byte", (PyCFunction)THPVariable_byte, METH_NOARGS, NULL},
   {"char", (PyCFunction)THPVariable_char, METH_NOARGS, NULL},
   {"contiguous", (PyCFunction)THPVariable_contiguous, METH_VARARGS | METH_KEYWORDS, NULL},

torch/_tensor_docs.py

@@ -749,7 +749,7 @@ def add_docstr_all(method, docstr):
     tensor([[5., 0., 0.],
             [0., 5., 0.],
             [0., 0., 5.]])
-    >>> b = torch.zeros(7, 3)        
+    >>> b = torch.zeros(7, 3)
     >>> b.fill_diagonal_(5)
     tensor([[5., 0., 0.],
             [0., 5., 0.],
@@ -2562,6 +2562,12 @@ def callable(a, b) -> number
 ``self.char()`` is equivalent to ``self.to(torch.int8)``. See :func:`to`.
 """)
 
+add_docstr_all('bfloat16',
+               r"""
+bfloat16() -> Tensor
+``self.bfloat16()`` is equivalent to ``self.to(torch.bfloat16)``. See :func:`to`.
+""")
+
 add_docstr_all('double',
                r"""
 double() -> Tensor