Merge pull request #175 from dsyme/funcs5

einsum

Author: dsyme

src/Native/LibTorchSharp/THSTensor.cpp

@@ -573,6 +573,11 @@ Tensor THSTensor_div_scalar_(const Tensor left, const Scalar right)
     CATCH_TENSOR(left->div_(*right));
 }
 
+Tensor THSTensor_einsum(const char* equation, const Tensor* tensors, const int length)
+{
+    CATCH_TENSOR(torch::einsum(equation, toTensors<at::Tensor>((torch::Tensor**)tensors, length)));
+}
+
 int64_t THSTensor_element_size(const Tensor tensor)
 {
     CATCH_RETURN(int64_t, 0, tensor->element_size());

src/Native/LibTorchSharp/THSTensor.h

@@ -215,6 +215,8 @@ EXPORT_API(Tensor) THSTensor_div_scalar(const Tensor left, const Scalar right);
 
 EXPORT_API(Tensor) THSTensor_div_scalar_(const Tensor left, const Scalar right);
 
+EXPORT_API(Tensor) THSTensor_einsum(const char* equation, const Tensor* tensors, const int length);
+
 EXPORT_API(int64_t) THSTensor_element_size(const Tensor tensor);
 
 EXPORT_API(Tensor) THSTensor_empty(

src/TorchSharp/Tensor/TorchTensor.cs

@@ -3746,23 +3746,39 @@ public static TorchTensor Cat(this TorchTensor[] tensors, long dimension)
                 return tensors[0];
             }
 
-            var parray = new PinnedArray<IntPtr>();
-            IntPtr tensorsRef = parray.CreateArray(tensors.Select(p => p.Handle).ToArray());
+            using (var parray = new PinnedArray<IntPtr>()) {
+                IntPtr tensorsRef = parray.CreateArray(tensors.Select(p => p.Handle).ToArray());
 
-            return new TorchTensor(THSTensor_cat(tensorsRef, parray.Array.Length, dimension));
+                return new TorchTensor(THSTensor_cat(tensorsRef, parray.Array.Length, dimension));
+            }
         }
 
         [DllImport("LibTorchSharp")]
         extern static IntPtr THSTensor_stack(IntPtr tensor, int len, long dim);
 
         public static TorchTensor Stack(this TorchTensor[] tensors, long dimension)
         {
-            var parray = new PinnedArray<IntPtr>();
-            IntPtr tensorsRef = parray.CreateArray(tensors.Select(p => p.Handle).ToArray());
+            using (var parray = new PinnedArray<IntPtr>()) {
+                IntPtr tensorsRef = parray.CreateArray(tensors.Select(p => p.Handle).ToArray());
 
-            var res = THSTensor_stack(tensorsRef, parray.Array.Length, dimension);
-            Torch.CheckForErrors();
-            return new TorchTensor(res);
+                var res = THSTensor_stack(tensorsRef, parray.Array.Length, dimension);
+                Torch.CheckForErrors();
+                return new TorchTensor(res);
+            }
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static IntPtr THSTensor_einsum([MarshalAs(UnmanagedType.LPStr)] string location, IntPtr tensors, int len);
+
+        public static TorchTensor Einsum(string equation, params TorchTensor[] tensors)
+        {
+            using (var parray = new PinnedArray<IntPtr>()) {
+                IntPtr tensorsRef = parray.CreateArray(tensors.Select(p => p.Handle).ToArray());
+
+                var res = THSTensor_einsum(equation, tensorsRef, parray.Array.Length);
+                Torch.CheckForErrors();
+                return new TorchTensor(res);
+            }
         }
 
     }