Implement tensor.size(int) for BatchedTensor

aten/src/ATen/BatchingRegistrations.cpp

@@ -69,6 +69,12 @@ TORCH_LIBRARY_IMPL(_, Batched, m) {
 }
 
 TORCH_LIBRARY_IMPL(aten, Batched, m) {
+  // NB: Ideally we would like some operators, like size.int, to "fallthrough"
+  // to the underlying implementation. However, because a BatchedTensor is a
+  // Tensor wrapper, it only has one dispatch key (Batched) on it. The resolution
+  // here is to just directly call the underlying implementation.
+  m.impl("size.int", static_cast<int64_t (*)(const Tensor&, int64_t)>(native::size));
+
   m.impl_UNBOXED("sum.dim_IntList", sum_batching_rule);
   m.impl_UNBOXED("mul.Tensor", mul_batching_rule);
 }

aten/src/ATen/test/vmap_test.cpp

@@ -612,5 +612,35 @@ TEST(VmapTest, TestBatchedTensorMul) {
   }
 }
 
+// test for BatchedTensor::size(int).
+TEST(VmapTest, TestBatchedTensorSize) {
+  {
+    // Single batch dim at front
+    Tensor x = at::randn({3, 5, 7});
+    Tensor Bx = makeBatched(x, {{0, 0}});
+
+    ASSERT_EQ(Bx.size(0), 5);
+    ASSERT_EQ(Bx.size(1), 7);
+    ASSERT_EQ(Bx.size(-1), 7);
+    ASSERT_EQ(Bx.size(-2), 5);
+    ASSERT_THROW(Bx.size(2), c10::Error);
+    ASSERT_THROW(Bx.size(-3), c10::Error);
+  }
+  {
+    // multiple batch dims not at front
+    Tensor x = at::randn({2, 3, 5, 7, 11});
+    Tensor Bx = makeBatched(x, {{0, 3}, {1, 1}});
+
+    ASSERT_EQ(Bx.size(0), 2);
+    ASSERT_EQ(Bx.size(1), 5);
+    ASSERT_EQ(Bx.size(2), 11);
+    ASSERT_EQ(Bx.size(-1), 11);
+    ASSERT_EQ(Bx.size(-2), 5);
+    ASSERT_EQ(Bx.size(-3), 2);
+    ASSERT_THROW(Bx.size(3), c10::Error);
+    ASSERT_THROW(Bx.size(-4), c10::Error);
+  }
+}
+
 
 }