move preparation

Author: Svetlana Karslioglu

prototype_source/maskedtensor_overview.py

@@ -20,7 +20,24 @@
 # 
 # Using MaskedTensor
 # ++++++++++++++++++
+# 
+# In this section we discuss how to use MaskedTensor including how to construct, access, the data
+# and mask, as well as indexing and slicing.
+#
+# Preparation
+# -----------
+# 
+# We'll begin by doing the necessary setup for the tutorial:
 #
+
+import torch
+from torch.masked import masked_tensor, as_masked_tensor
+import warnings
+
+# Disable prototype warnings and such
+warnings.filterwarnings(action='ignore', category=UserWarning)
+
+######################################################################
 # Construction
 # ------------
 # 
@@ -52,9 +69,6 @@
 # as :class:`torch.Tensor`. Below are some examples of common indexing and slicing patterns:
 #
 
-import torch
-from torch.masked import masked_tensor, as_masked_tensor
-
 data = torch.arange(24).reshape(2, 3, 4)
 mask = data % 2 == 0
 
@@ -174,7 +188,6 @@
 x = torch.tensor([1., 1.], requires_grad=True)
 div = torch.tensor([0., 1.])
 y = x/div # => y is [inf, 1]
-    >>>
 mask = (div != 0) # => mask is [0, 1]
 loss = as_masked_tensor(y, mask)
 loss.sum().backward()
@@ -213,7 +226,7 @@
 # Safe Softmax
 # ------------
 #
-# Safe softmax is another great example of `an issue <https://github.com/pytorch/pytorch/issues/55056>`_
+# Safe softmax is another great example of `an issue <https://github.com/pytorch/pytorch/issues/55056>`__
 # that arises frequently. In a nutshell, if there is an entire batch that is "masked out"
 # or consists entirely of padding (which, in the softmax case, translates to being set `-inf`),
 # then this will result in NaNs, which can lead to training divergence.
@@ -247,15 +260,16 @@
 
 ######################################################################
 # Implementing missing torch.nan* operators
-# --------------------------------------------------------------------------------------------------------------
+# -----------------------------------------
 #
 # In `Issue 61474 <<https://github.com/pytorch/pytorch/issues/61474>`__,
 # there is a request to add additional operators to cover the various `torch.nan*` applications,
 # such as ``torch.nanmax``, ``torch.nanmin``, etc.
 #
 # In general, these problems lend themselves more naturally to masked semantics, so instead of introducing additional
-# operators, we propose using :class:`MaskedTensor`s instead. Since
-# `nanmean has already landed <https://github.com/pytorch/pytorch/issues/21987>`_, we can use it as a comparison point:
+# operators, we propose using :class:`MaskedTensor`s instead.
+# Since `nanmean has already landed <https://github.com/pytorch/pytorch/issues/21987>`__,
+# we can use it as a comparison point:
 #
 
 x = torch.arange(16).float()