Add torch.nested namespace (pytorch#84102)

First step towards pytorch#83775
- only `to_padded_tensor` is moved to the nested namespace for now
- following the schema used for `special`, `fft`, `linalg` and other namespaces, nested functions are registered in native_functions.yaml as `nested_{function_name}` and are bound to the desired Python name in
`torch/nested/__init__.py`, and the desired C++ name in `torch/csrc/api/include/torch/nested.h`.

~~**Question**: should we keep the documentation for `Tensor.to_padded_tensor` or can this deleted since it is shared by `torch.nested.to_padded_tensor`?~~

[generated nested docs](https://docs-preview.pytorch.org/84102/nested.html?highlight=nested#module-torch.nested)

Differential Revision: [D39361148](https://our.internmc.facebook.com/intern/diff/D39361148)
Pull Request resolved: pytorch#84102
Approved by: https://github.com/drisspg

Author: mikaylagawarecki

aten/src/ATen/native/README.md

@@ -628,7 +628,8 @@ the torch._C._nn (marked with `python_module: nn`),
 torch._C._fft (marked with `python_module: fft`),
 torch._C._linalg (marked with `python_module: linalg`) objects,
 torch._C._sparse (marked with `python_module: sparse`) objects,
-or torch._C._special (marked with `python_module: special`) objects.
+torch._C._special (marked with `python_module: special`) objects,
+or torch._C._nested (marked with `python_module: nested`) objects.
 
 ### Undefined tensor conventions
 

aten/src/ATen/native/native_functions.yaml

@@ -12558,6 +12558,21 @@
 - func: linalg_multi_dot.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!)
   python_module: linalg
 
+## Functions related to the `torch.nested` namespace
+# Note [nested namespace binding]
+# Functions in the nested python module should have their names start with
+#   "nested_" underscore and be bound to the desired Python name in
+#   torch/nested/__init__.py, and the desired C++ name in torch/csrc/api/include/torch/nested.h.
+#   The "nested_" names should be hidden from the user and not documented.
+
+- func: nested_to_padded_tensor(Tensor self, float padding, int[]? output_size=None) -> Tensor
+  python_module: nested
+  variants: function
+  dispatch:
+    NestedTensorCPU: NestedTensor_to_padded_tensor_generic
+    NestedTensorCUDA: NestedTensor_to_padded_tensor_cuda
+  autogen: nested_to_padded_tensor.out
+
 ## Functions that are only for testing
 # It is undocumented and should not be used outside of tests.
 - func: _test_serialization_subcmul(Tensor self, Tensor other, Scalar alpha=1) -> Tensor
@@ -13068,13 +13083,6 @@
   dispatch:
     CompositeExplicitAutograd: alias_copy_out
 
-- func: to_padded_tensor(Tensor self, float padding, int[]? output_size=None) -> Tensor
-  variants: method
-  dispatch:
-    NestedTensorCPU: NestedTensor_to_padded_tensor_generic
-    NestedTensorCUDA: NestedTensor_to_padded_tensor_cuda
-  autogen: to_padded_tensor.out
-
 - func: _nested_tensor_softmax_with_shape(Tensor self, Tensor query) -> Tensor
   dispatch:
     NestedTensorCPU: NestedTensor_softmax_dropout

aten/src/ATen/native/transformers/attention.cpp

@@ -262,7 +262,7 @@ std::tuple<Tensor, Tensor, Tensor> transform_bias_rescale_qkv_cpu(
     const Tensor& qkv_bias,
     const int64_t num_head) {
   auto qkv_ = qkv.is_nested()
-    ? c10::MaybeOwned<Tensor>::owned(qkv.to_padded_tensor(0))
+    ? c10::MaybeOwned<Tensor>::owned(nested_to_padded_tensor(qkv, 0))
     : c10::MaybeOwned<Tensor>::borrowed(qkv);
   auto B = qkv_->size(0);
   auto T = qkv_->size(1);

build.bzl

@@ -257,6 +257,7 @@ _GENERATED_AUTOGRAD_PYTHON_CPP = [
     "torch/csrc/autograd/generated/python_functions_3.cpp",
     "torch/csrc/autograd/generated/python_functions_4.cpp",
     "torch/csrc/autograd/generated/python_nn_functions.cpp",
+    "torch/csrc/autograd/generated/python_nested_functions.cpp",
     "torch/csrc/autograd/generated/python_fft_functions.cpp",
     "torch/csrc/autograd/generated/python_linalg_functions.cpp",
     "torch/csrc/autograd/generated/python_return_types.cpp",

build_variables.bzl

@@ -986,6 +986,7 @@ def glob_libtorch_python_sources(gencode_pattern = ":generate-code[{}]"):
         "torch/csrc/autograd/generated/python_functions_2.cpp",
         "torch/csrc/autograd/generated/python_functions_3.cpp",
         "torch/csrc/autograd/generated/python_functions_4.cpp",
+        "torch/csrc/autograd/generated/python_nested_functions.cpp",
         "torch/csrc/autograd/generated/python_nn_functions.cpp",
         "torch/csrc/autograd/generated/python_fft_functions.cpp",
         "torch/csrc/autograd/generated/python_linalg_functions.cpp",

caffe2/CMakeLists.txt

@@ -395,6 +395,7 @@ set(GENERATED_CXX_PYTHON
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_nn_functions.cpp"
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_fft_functions.cpp"
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_linalg_functions.cpp"
+  "${TORCH_SRC_DIR}/csrc/autograd/generated/python_nested_functions.cpp"
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_sparse_functions.cpp"
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_special_functions.cpp"
   "${TORCH_SRC_DIR}/csrc/autograd/generated/python_return_types.cpp"

docs/source/index.rst

@@ -92,7 +92,7 @@ Features described in this documentation are classified by release status:
    rpc
    torch.random <random>
    masked
-   nested
+   torch.nested <nested>
    sparse
    storage
    torch.testing <testing>

docs/source/nested.rst

@@ -142,14 +142,11 @@ nested_tensor([
 
 Note that nt.unbind()[0] is not a, but rather a slice of the underlying memory, which represents the first entry or constituent of the NestedTensor.
 
-Nested tensor methods
+Nested tensor functions
 +++++++++++++++++++++++++
 
-The following Tensor methods are related to nested tensors:
+The following functions are related to nested tensors:
 
-.. currentmodule:: torch
-.. autosummary::
-    :toctree: generated
-    :nosignatures:
+.. currentmodule:: torch.nested
 
-    Tensor.to_padded_tensor
+.. autofunction:: to_padded_tensor

test/allowlist_for_publicAPI.json

@@ -1216,6 +1216,9 @@
   "torch.multiprocessing.spawn": [
     "Optional"
   ],
+  "torch.nested": [
+    "to_padded_tensor"
+  ],
   "torch.nn.common_types": [
     "Optional",
     "Tensor",

test/cpp/api/CMakeLists.txt

@@ -18,6 +18,7 @@ set(TORCH_API_TEST_SOURCES
   ${TORCH_API_TEST_DIR}/moduledict.cpp
   ${TORCH_API_TEST_DIR}/modulelist.cpp
   ${TORCH_API_TEST_DIR}/modules.cpp
+  ${TORCH_API_TEST_DIR}/nested.cpp
   ${TORCH_API_TEST_DIR}/parameterdict.cpp
   ${TORCH_API_TEST_DIR}/parameterlist.cpp
   ${TORCH_API_TEST_DIR}/namespace.cpp

test/cpp/api/nested.cpp

@@ -0,0 +1,15 @@
+#include <gtest/gtest.h>
+
+#include <torch/nested.h>
+#include <torch/torch.h>
+
+#include <test/cpp/api/support.h>
+
+// Simple test that verifies the nested namespace is registered properly
+//   properly in C++
+TEST(NestedTest, Nested) {
+  auto a = torch::randn({2, 3});
+  auto b = torch::randn({4, 5});
+  auto nt = torch::nested_tensor({a, b});
+  torch::nested::to_padded_tensor(nt, 0);
+}

test/forward_backward_compatibility/check_forward_backward_compatibility.py

@@ -278,6 +278,7 @@
     # Distributed c10d ops are all going to be updated
     ("c10d::.*", datetime.date(2022, 10, 31)),
     ("c10d::allgather_", datetime.date(2022, 10, 1)),
+    ("aten::to_padded_tensor", datetime.date(2022, 10, 1)),
 ]
 
 ALLOW_LIST_COMPILED = [

test/test_native_mha.py

@@ -213,7 +213,7 @@ def forward(self, q, k, v, key_padding_mask):
             q, k, v, key_padding_mask=mask if use_padding and not use_nt else None
         )
         if use_nt:
-            ynpt = ynpt.to_padded_tensor(0)
+            ynpt = torch.nested.to_padded_tensor(ynpt, 0)
             if pad_all:
                 ynpt_final = torch.zeros_like(ypt)
                 ynpt_final[:, :ynpt.shape[1], :] = ynpt