Check operator coverage (forward only) for issue 313 (#316)

Author: cpuhrsch

nestedtensor/csrc/BinaryOps.cpp

@@ -274,6 +274,10 @@ TORCH_LIBRARY_IMPL(aten, NestedTensor, m) {
   nt_impl(m, "logical_and_", NestedTensor_binary_<at::native::logical_and_>);
   nt_impl(m, "logical_and.out", NestedTensor_binary_out<at::logical_and_out>);
 
+  nt_impl(m, "logical_or", NestedTensor_binary<at::logical_or>);
+  nt_impl(m, "logical_or_", NestedTensor_binary_<at::native::logical_or_>);
+  nt_impl(m, "logical_or.out", NestedTensor_binary_out<at::logical_or_out>);
+
   nt_impl(m, "sub.Tensor", (NestedTensor_binary<Scalar, at::sub>));
   nt_impl(m, "pow.Tensor_Tensor", NestedTensor_binary<at::pow>);
 }

nestedtensor/csrc/ReduceOps.cpp

@@ -102,6 +102,50 @@ Tensor NestedTensor_sum_dim(
       my_sum, self, dims, keepdims, dtype);
 }
 
+std::tuple<Tensor, Tensor> NestedTensor_max_dim(
+    const Tensor& self,
+    int64_t dim,
+    bool keepdims) {
+  int64_t nested_dim = get_nested_tensor_impl(self)->nested_dim();
+  at::Tensor output = self;
+  if (dim >= nested_dim) {
+    std::vector<TensorNode> result = unzip(map(
+        [nested_dim, dim, keepdims](at::Tensor tensor) {
+          auto tmp = at::max(tensor, dim - nested_dim, keepdims);
+          std::vector<at::Tensor> result;
+          result.push_back(std::get<0>(tmp));
+          result.push_back(std::get<1>(tmp));
+          return result;
+        },
+        get_nested_tensor_structure(output)));
+    return std::make_tuple(
+        wrap_tensor_node(std::move(result[0])),
+        wrap_tensor_node(std::move(result[1])));
+  }
+  auto opt_sizes = get_opt_sizes(output);
+  TORCH_CHECK(
+      opt_sizes[dim],
+      "Current shape doesn't support reduction across nested dimension. Please open a feature request https://t.ly/62F6.");
+  auto new_nested_size = get_nested_size(output);
+  new_nested_size = squeeze(new_nested_size, dim, keepdims);
+  auto tmp =
+      at::max(NestedTensor_to_tensor(output, c10::nullopt), dim, keepdims);
+  return std::make_tuple(
+      wrap_buffer(std::get<0>(tmp).reshape({-1}), new_nested_size),
+      wrap_buffer(std::get<1>(tmp).reshape({-1}), new_nested_size));
+}
+
+Tensor NestedTensor_max(const Tensor& self) {
+  auto tensors = flatten(
+      map([](at::Tensor tensor) { return at::max(tensor); },
+          get_nested_tensor_structure(self)));
+  if (tensors.size() == 0) {
+    return at::ones({0});
+  }
+  auto all_tensor = at::stack(tensors);
+  return at::max(all_tensor);
+}
+
 Tensor NestedTensor_mean_dim(
     const Tensor& self,
     c10::ArrayRef<int64_t> dims,
@@ -336,6 +380,8 @@ TORCH_LIBRARY_IMPL(aten, NestedTensor, m) {
   nt_impl(m, "sum.dim_IntList", NestedTensor_sum_dim);
   nt_impl(m, "mean", NestedTensor_mean);
   nt_impl(m, "mean.dim", NestedTensor_mean_dim);
+  nt_impl(m, "max", NestedTensor_max);
+  nt_impl(m, "max.dim", NestedTensor_max_dim);
   nt_impl(m, "var", NestedTensor_var);
   nt_impl(m, "var.dim", NestedTensor_var_dim);
   nt_impl(m, "var_backward.dim", NestedTensor_var_backward_dim);

nestedtensor/version.py

@@ -1,5 +1,5 @@
-__version__ = '0.0.1+2c3a468'
-git_version = '2c3a468de0a4a2e8c50d6dd7b41282fe98471206'
+__version__ = '0.0.1+6a95af1'
+git_version = '6a95af1cfc7efd07c7d103b3b3e5cd27148b128f'
 from nestedtensor import _C
 if hasattr(_C, 'CUDA_VERSION'):
     cuda = _C.CUDA_VERSION

test/test_coverage.py

@@ -0,0 +1,47 @@
+import traceback
+import functools
+import pdb
+import sys
+import torch
+import nestedtensor
+import unittest
+import random
+from torch.nn import functional as F
+from torch import nn
+
+from utils import TestCase
+
+
+def ntnt(x): return nestedtensor.nested_tensor(x, requires_grad=True)
+def ntnt_nograd(x): return nestedtensor.nested_tensor(x)
+
+
+# Various smoke tests to confirm coverage of an operator
+
+class TestCoverage(TestCase):
+
+    def test_issues_313(self):
+        # Based on https://github.com/pytorch/nestedtensor/issues/313
+
+        def model(x):
+            torch.manual_seed(20)
+            linear = nn.Linear(9, 64)
+            norm = nn.BatchNorm1d(64)
+            # 3 voxel with 40, 50 and 90 points respectively
+            x = linear(x)
+            x = norm(x.transpose(2, 1).contiguous()
+                     ).transpose(2, 1).contiguous()
+            x = F.relu(x)
+            return torch.max(x, dim=1, keepdim=True)[0]
+
+        inputs = [torch.randn(i, 9) for i in [40, 50, 90]]
+        model(ntnt(inputs))
+
+        inputs = [torch.randn(30, 9) for _ in range(3)]
+        x0 = model(ntnt(inputs))
+        x1 = model(torch.stack(inputs))
+        self.assertEqual(torch.stack(x0.unbind()), x1)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/test_nested_tensor_reduce.py

@@ -30,13 +30,13 @@ def _flatten_nt(nt):
 
 class TestReduce(TestCase):
 
-    def _test_reduce_dim(self, fn, associative=True, test_keep_dim=True):
+    def _test_reduce_dim(self, fn, associative=True, test_keep_dim=True, test_multi_dim=True):
         t0 = torch.arange(9).float().reshape(3, 3)
         t1 = torch.arange(6).float().reshape(2, 3)
         t2 = torch.arange(9).float().reshape(3, 3)
         ts = [[t0, t1], [t2, t1]]
         nt = ntnt(ts)
-        if associative:
+        if associative and test_multi_dim:
             t01 = fn(torch.stack([fn(t0, 0), fn(t1, 0)]), 0)
             t21 = fn(torch.stack([fn(t2, 0), fn(t1, 0)]), 0)
             t02 = fn(torch.stack([fn(t0, 0), fn(t2, 0)]), 0)
@@ -125,6 +125,13 @@ def test_sum_all(self):
     def test_sum_dim(self):
         self._test_reduce_dim(torch.sum, True)
 
+    def test_max_all(self):
+        self._test_allreduce(lambda x: x.max())
+
+    def test_max_dim(self):
+        self._test_reduce_dim(lambda x, dim, keepdim=False: x.max(
+            dim, keepdim)[0], True, test_multi_dim=False)
+
     def test_mean_all(self):
         self._test_allreduce(lambda x: x.mean())