Add Tensor.Op.Reshape.

spec/Main.savi

@@ -10,6 +10,7 @@
       Spec.Run(Tensor.Op.Lesser.Spec).new(env)
       Spec.Run(Tensor.Op.Logical.Spec).new(env)
       Spec.Run(Tensor.Op.MatMul.Spec).new(env)
+      Spec.Run(Tensor.Op.Reshape.Spec).new(env)
       Spec.Run(Tensor.Op.Select.Spec).new(env)
       Spec.Run(Tensor.Op.Softmax.Spec).new(env)
     ])

spec/Tensor.Op.Reshape.Spec.savi

@@ -0,0 +1,24 @@
+:class Tensor.Op.Reshape.Spec
+  :is Spec
+  :const describes: "Tensor.Op.Reshape"
+
+  :it "emits a variation on the tensor which has a changed shape"
+    _WithGraphHelper.run(@env) -> (g, session | assert no_error: (
+      result = session.compute!(
+        g.reshape!("example"
+          g.const!("input", Tensor(F64).from_array([1, 2, 3, 4, 5, 6]))
+          [2, 3]
+        )
+      )
+
+      assert: result.as!(Tensor(F64)).into_array == [1, 2, 3, 4, 5, 6]
+      assert: result.shape_into_array == [2, 3]
+    ))
+
+  :it "complains when the requested shape doesn't align with the current size"
+    _WithGraphHelper.run(@env, False) -> (g, session |
+      assert error: g.reshape!("example"
+        g.const!("input", Tensor(F64).from_array([1, 2, 3, 4, 5, 6, 7, 8]))
+        [2, 3]
+      )
+    )

src/Tensor.Graph.Helper.savi

@@ -40,7 +40,7 @@
     Tensor.Op.Lesser.new!(@graph, name, x, y, True)
 
   ///
-  // Type Conversions
+  // Type/Shape Conversions
 
   :fun ref bitcast!(name, input, output_type)
     Tensor.Op.Bitcast.new!(@graph, name, input, output_type)
@@ -51,6 +51,15 @@
   :fun ref cast_with_floating_point_truncation!(name, input, output_type)
     Tensor.Op.Cast.new!(@graph, name, input, output_type, True)
 
+  :fun ref reshape!(name, input, output_shape Array(USize))
+    Tensor.Op.Reshape.new!(@graph, name, input
+      @const!("\(name).new_shape"
+        Tensor(I64).generate(output_shape.size) -> (index |
+          output_shape[index]!.i64!
+        )
+      )
+    )
+
   ///
   // Other Unary Operations
 

src/Tensor.Op.Reshape.savi

@@ -0,0 +1,17 @@
+:: Create a new output tensor of a different shape using the elements copied
+:: from the input tensor. The second input tensor indicates the new shape.
+::
+:: It is similar to the `Tensor.reshape` method, but it is used for runtime
+:: values inside the graph rather than tensors held locally (not in a graph).
+:struct box Tensor.Op.Reshape
+  :is Tensor.Op
+  :fun non new!(graph Tensor.Graph, name
+    input
+    output_shape Tensor.Graph.CanOutput
+  )
+    @_new(graph.new_operation("Reshape", name) -> (builder |
+      builder
+        .add_input(input)
+        .add_input(output_shape)
+        .finish!
+    ))

src/Tensor.savi

@@ -1,6 +1,10 @@
 :trait Tensor.Any
   :let _ptr CPointer(_FFI.Tensor)
   :let _ptr_is_owned Bool
+
+  :fun shape_into_array(shape Array(USize) = []) Array(USize)
+  :fun element_count USize
+  :fun element_byte_width U8
   :fun non element_type_code I32: -1
 
 :class Tensor(T Numeric(T)'val) // TODO: Exclude non-machine word implementers of the Numeric trait
@@ -62,6 +66,15 @@
       total_byte_size
     )
 
+  :fun non generate(total_element_count USize)
+    :yields USize for T
+    // TODO: Do this without the intermediate array?
+    elements Array(T) = []
+    total_element_count.times -> (index |
+      elements << yield index
+    )
+    @from_array(elements)
+
   :fun ref try_reshape(dimensions): try (@reshape!(dimensions) | @)
   :fun ref reshape!(dimensions Array(USize))
     element_count USize = 1
@@ -103,6 +116,12 @@
     )
     data
 
+  :fun shape_into_array(shape Array(USize) = []) Array(USize)
+    num_dims = @_ffi.num_dims(@_ptr).usize
+    shape.reserve(shape.size + num_dims)
+    num_dims.times -> (index | shape << @_ffi.dim_at(@_ptr, index.i32).usize)
+    shape
+
   :fun element_count: _FFI.Tensor.element_count(@_ptr).usize
   :fun element_byte_width: T.byte_width
   :fun non element_type_code I32: _FFI.DataType(T).code // TODO: avoid I32 here?
@@ -131,6 +150,12 @@
   :ffi element_count(tensor CPointer(@)) U64
     :foreign_name TF_TensorElementCount
 
+  :ffi num_dims(tensor CPointer(@)) I32
+    :foreign_name TF_NumDims
+
+  :ffi dim_at(tensor CPointer(@), index I32) I64
+    :foreign_name TF_Dim
+
   :ffi set_shape(
     tensor CPointer(@)
     dimension_list CPointer(I64)