wip add llama

sbrunk · Jul 25, 2023 · 91452c1 · 91452c1
1 parent 938c7bf
commit 91452c1
Show file tree

Hide file tree

Showing 2 changed files with 385 additions and 1 deletion.
diff --git a/core/src/main/scala/torch/Tensor.scala b/core/src/main/scala/torch/Tensor.scala
@@ -334,7 +334,7 @@ sealed abstract class Tensor[D <: DType]( /* private[torch]  */ val native: pyto
     *   x.expand(-1, 4) // -1 means not changing the size of that dimension
     *   ```
     */
-  def expand(sizes: Int*) = Tensor(native.expand(sizes.map(_.toLong)*))
+  def expand(sizes: Int*): Tensor[D] = Tensor(native.expand(sizes.map(_.toLong)*))
 
   def flatten: Tensor[D] = Tensor(native.flatten())
 
@@ -545,6 +545,11 @@ sealed abstract class Tensor[D <: DType]( /* private[torch]  */ val native: pyto
     */
   def t: Tensor[D] = Tensor(native.t())
 
+  /** Returns a tensor that is a transposed version of this tensor. The given dimensions dim0 and
+    * dim1 are swapped.
+    */
+  def transpose(dim0: Int, dim1: Int): Tensor[D] = Tensor(native.transpose(dim0, dim1))
+
   /** Calculates the variance of all elements of this tensor. */
   def variance = Tensor(native.`var`())