pytorch
diff --git a/‎test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp‎
Lines changed: 18 additions & 4 deletions b/‎test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp‎
Lines changed: 18 additions & 4 deletions
diff --git a/‎test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/ops.py‎
Lines changed: 12 additions & 0 deletions b/‎test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/ops.py‎
Lines changed: 12 additions & 0 deletions
diff --git a/‎test/cpp_extensions/libtorch_agnostic_extension/test/test_libtorch_agnostic.py‎
Lines changed: 7 additions & 0 deletions b/‎test/cpp_extensions/libtorch_agnostic_extension/test/test_libtorch_agnostic.py‎
Lines changed: 7 additions & 0 deletions
@@ -355,32 +355,46 @@ Tensor my_zero_(Tensor t) {
   return zero_(t);
 }
 
-Tensor my_amax(Tensor t) {
-  return amax(t, 0, false);
-}
-
 void boxed_my_zero_(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
   auto res = my_zero_(to<Tensor>(stack[0]));
   stack[0] = from(res);
 }
 
+Tensor my_amax(Tensor t) {
+  return amax(t, 0, false);
+}
+
 void boxed_my_amax(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
   auto res = my_amax(to<Tensor>(stack[0]));
   stack[0] = from(res);
 }
 
+Tensor my_amax_vec(Tensor t) {
+  std::vector<int64_t> v = {0};
+  return amax(t, v, false);
+}
+
+void boxed_my_amax_vec(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
+  auto res = my_amax_vec(to<Tensor>(stack[0]));
+  stack[0] = from(res);
+}
+
+
 STABLE_TORCH_LIBRARY_FRAGMENT(libtorch_agnostic, m) {
   m.def("my_zero_(Tensor(a!) t) -> Tensor(a!)");
   m.def("my_amax(Tensor a) -> Tensor");
+  m.def("my_amax_vec(Tensor a) -> Tensor");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(libtorch_agnostic, CPU, m) {
   m.impl("my_zero_", &boxed_my_zero_);
   m.impl("my_amax", &boxed_my_amax);
+  m.impl("my_amax_vec", &boxed_my_amax_vec);
 }
 
 STABLE_TORCH_LIBRARY_IMPL(libtorch_agnostic, CUDA, m) {
   m.impl("my_amax", &boxed_my_amax);
+  m.impl("my_amax_vec", &boxed_my_amax_vec);
 }
 
 bool test_default_constructor(bool defined) {
 
@@ -166,6 +166,18 @@ def my_amax(t) -> Tensor:
     return torch.ops.libtorch_agnostic.my_amax.default(t)
 
 
+def my_amax_vec(t) -> Tensor:
+    """
+    Returns t.amax()
+
+    Args:
+        t: Tensor
+
+    Returns: amax(t)
+    """
+    return torch.ops.libtorch_agnostic.my_amax_vec.default(t)
+
+
 def fill_infinity(t) -> Tensor:
     """
     Fills the tensor with inf.
 
@@ -216,6 +216,13 @@ def test_my_amax(self, device):
             out = libtorch_agnostic.ops.my_amax(t)
             self.assertEqual(out, torch.amax(t, 0))
 
+        def test_my_amax_vec(self, device):
+            import libtorch_agnostic
+
+            t = torch.rand(2, 7, device=device)
+            out = libtorch_agnostic.ops.my_amax_vec(t)
+            self.assertEqual(out, torch.amax(t, 0))
+
         def test_fill_infinity(self, device):
             import libtorch_agnostic