From 1bea5fc3ad565c91886244be259706efffa99f9e Mon Sep 17 00:00:00 2001
From: Viswanath Sivakumar <viswanath@fb.com>
Date: Wed, 24 Oct 2018 13:08:45 -0700
Subject: [PATCH] Fix UpsampleNearest op CPU impl batch handling (#13002)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13002

Batch dim wasn't handled in the CPU impl (will fail for inputs with N > 1).
Fixing that here.

Differential Revision: D10515159

fbshipit-source-id: ee7e4f489d2d4de793f550b31db7c0e2ba3651e8
---
 modules/detectron/upsample_nearest_op.h       |  2 +-
 modules/detectron/upsample_nearest_op_test.py | 43 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 modules/detectron/upsample_nearest_op_test.py

diff --git a/modules/detectron/upsample_nearest_op.h b/modules/detectron/upsample_nearest_op.h
index ba5890400a9d5..a891de420cb9d 100644
--- a/modules/detectron/upsample_nearest_op.h
+++ b/modules/detectron/upsample_nearest_op.h
@@ -50,7 +50,7 @@ class UpsampleNearestOp final : public Operator<Context> {
       d2 = Y->dim32(1);
       d3 = Y->dim32(2);
     } else {
-      d1 = Y->dim32(1);
+      d1 = Y->dim32(0) * Y->dim32(1);
       d2 = Y->dim32(2);
       d3 = Y->dim32(3);
     }
diff --git a/modules/detectron/upsample_nearest_op_test.py b/modules/detectron/upsample_nearest_op_test.py
new file mode 100644
index 0000000000000..5eef675029ab9
--- /dev/null
+++ b/modules/detectron/upsample_nearest_op_test.py
@@ -0,0 +1,43 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import unittest
+
+import caffe2.python.hypothesis_test_util as hu
+import hypothesis.strategies as st
+import numpy as np
+from caffe2.python import core, dyndep
+from hypothesis import given
+
+
+dyndep.InitOpsLibrary("@/caffe2/modules/detectron:detectron_ops")
+
+
+class TestUpsampleNearestOp(hu.HypothesisTestCase):
+    @given(
+        N=st.integers(1, 3),
+        H=st.integers(10, 300),
+        W=st.integers(10, 300),
+        scale=st.integers(1, 3),
+        **hu.gcs
+    )
+    def test_upsample_nearest_op(self, N, H, W, scale, gc, dc):
+        C = 32
+        X = np.random.randn(N, C, H, W).astype(np.float32)
+        op = core.CreateOperator("UpsampleNearest", ["X"], ["Y"], scale=scale)
+
+        def ref(X):
+            outH = H * scale
+            outW = W * scale
+            outH_idxs, outW_idxs = np.meshgrid(
+                np.arange(outH), np.arange(outW), indexing="ij"
+            )
+            inH_idxs = (outH_idxs / scale).astype(np.int32)
+            inW_idxs = (outW_idxs / scale).astype(np.int32)
+            Y = X[:, :, inH_idxs, inW_idxs]
+            return [Y]
+
+        self.assertReferenceChecks(device_option=gc, op=op, inputs=[X], reference=ref)
+
+
+if __name__ == "__main__":
+    unittest.main()