Fix Rust codegen matrix return types

symforce/codegen/backends/rust/templates/util/util.jinja

@@ -59,7 +59,7 @@
         {% if is_symbolic(type) or T.__name__ == "float" %}
             {{ format_scalar(spec) }}
         {% elif issubclass(T, Matrix) %}
-            {{ format_vector(T, spec.return_key, spec) }}
+            {{ format_typename(spec.return_key, type, spec) }}
         {% else %}
             {{ raise("Unsupported return type: {}".format(T)) }}
         {% endif %}
@@ -231,8 +231,7 @@
             {% if is_symbolic(type) or T_return.__name__ == "float" %}
             {{ terms[0][1] }};
             {% elif issubclass(T_return, Matrix) %}
-                {% set size = T_return.SHAPE[0] * T_return.SHAPE[1] %}
-                {{ format_vector(T_return, name, spec) }}::new(
+                {{ format_typename(name, type, spec) }}::new(
                 {% for lhs, rhs in terms %}
         {{ rhs }}{% if not loop.last %},{% endif %}
                 {% endfor %}

test/symforce_function_codegen_test_data/symengine/symforce_rust_codegen_test/src/lib.rs

@@ -1,4 +1,5 @@
 
 mod backend_test_function_float32;
 mod backend_test_function_float64;
+mod matrix_return_fun;
 mod vector_matrix_fun;

test/symforce_function_codegen_test_data/symengine/symforce_rust_codegen_test/src/matrix_return_fun.rs

@@ -0,0 +1,28 @@
+// -----------------------------------------------------------------------------
+// This file was autogenerated by symforce from template:
+//     function/FUNCTION.rs.jinja
+// Do NOT modify by hand.
+// -----------------------------------------------------------------------------
+
+pub mod sym {
+
+    #[allow(unused_parens)]
+
+    ///
+    /// This function was autogenerated from a symbolic function. Do not modify by hand.
+    ///
+    /// Symbolic function: matrix_return
+    ///
+    /// Outputs:
+    ///     res: Matrix23
+
+    pub fn matrix_return_fun() -> nalgebra::SMatrix<f64, 2, 3> {
+        // Total ops: 0
+
+        // Intermediate terms (0)
+
+        // Output terms (1)
+
+        nalgebra::SMatrix::<f64, 2, 3>::new(1_f64, 4_f64, 2_f64, 5_f64, 3_f64, 6_f64)
+    }
+} // mod sym

test/symforce_function_codegen_test_data/sympy/symforce_rust_codegen_test/src/lib.rs

@@ -1,4 +1,5 @@
 
 mod backend_test_function_float32;
 mod backend_test_function_float64;
+mod matrix_return_fun;
 mod vector_matrix_fun;

test/symforce_function_codegen_test_data/sympy/symforce_rust_codegen_test/src/matrix_return_fun.rs

@@ -0,0 +1,28 @@
+// -----------------------------------------------------------------------------
+// This file was autogenerated by symforce from template:
+//     function/FUNCTION.rs.jinja
+// Do NOT modify by hand.
+// -----------------------------------------------------------------------------
+
+pub mod sym {
+
+    #[allow(unused_parens)]
+
+    ///
+    /// This function was autogenerated from a symbolic function. Do not modify by hand.
+    ///
+    /// Symbolic function: matrix_return
+    ///
+    /// Outputs:
+    ///     res: Matrix23
+
+    pub fn matrix_return_fun() -> nalgebra::SMatrix<f64, 2, 3> {
+        // Total ops: 0
+
+        // Intermediate terms (0)
+
+        // Output terms (1)
+
+        nalgebra::SMatrix::<f64, 2, 3>::new(1_f64, 4_f64, 2_f64, 5_f64, 3_f64, 6_f64)
+    }
+} // mod sym

test/symforce_rust_codegen_test.py

@@ -55,6 +55,7 @@ def write_lib_rs(output_dir: Path) -> None:
             """
             mod backend_test_function_float32;
             mod backend_test_function_float64;
+            mod matrix_return_fun;
             mod vector_matrix_fun;
             """
         )
@@ -91,6 +92,9 @@ def test_codegen(self) -> None:
         def rust_func(vec3: sf.V3, mat33: sf.M33) -> sf.Matrix31:
             return sf.Matrix31(mat33 * vec3)
 
+        def matrix_return() -> sf.M23:
+            return sf.M23([[1, 2, 3], [4, 5, 6]])
+
         output_dir_base = self.make_output_dir("symforce_rust_codegen_test_")
         output_dir_src = output_dir_base / "src"
 
@@ -103,6 +107,13 @@ def rust_func(vec3: sf.V3, mat33: sf.M33) -> sf.Matrix31:
             name="vector_matrix_fun",
         ).generate_function(output_dir_src, skip_directory_nesting=True)
 
+        # A non-square matrix distinguishes its row and column dimensions from a flattened vector.
+        Codegen.function(
+            matrix_return,
+            config=RustConfig(scalar_type=ScalarType.DOUBLE),
+            name="matrix_return_fun",
+        ).generate_function(output_dir_src, skip_directory_nesting=True)
+
         # Generate the symbolic backend test function
         for scalar in scalars:
             Codegen.function(