Rust codegen uses wrong dimensions for return types

[avsaase]

Describe the bug When generating Rust code from a function that returns a matrix, the generated function returns a vector with the same number of elements as the matrix.

To Reproduce Steps to reproduce the behavior, e.g.:

Take this script:

import symforce

symforce.set_epsilon_to_number()
import symforce.symbolic as sf
from symforce import codegen
from symforce.codegen.backends.rust.rust_config import RustConfig
from symforce.notebook_util import display_code_file


def create_matrix() -> sf.Matrix:
    return sf.Matrix55(
        1, 2, 3, 4, 5,
        6, 7, 8, 9, 10,
        11, 12, 13, 14, 15,
        16, 17, 18, 19, 20,
        21, 22, 23, 24, 25,
    )


output_dir = "./generated"
config = RustConfig()

codegen = codegen.Codegen.function(
    create_matrix,
    config=RustConfig(),
)
data = codegen.generate_function(output_dir)

This generates the following Rust file

// -----------------------------------------------------------------------------
// 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: create_matrix
    ///
    /// Outputs:
    ///     res: Matrix55

    pub fn create_matrix() -> nalgebra::SVector<f64, 25> {
        // Total ops: 0

        // Intermediate terms (0)

        // Output terms (1)

        nalgebra::SVector::<f64, 25>::new(
            1_f64, 6_f64, 11_f64, 16_f64, 21_f64, 2_f64, 7_f64, 12_f64, 17_f64, 22_f64, 3_f64,
            8_f64, 13_f64, 18_f64, 23_f64, 4_f64, 9_f64, 14_f64, 19_f64, 24_f64, 5_f64, 10_f64,
            15_f64, 20_f64, 25_f64,
        )
    }
} // mod sym

As you can see the generated function returns a nalgebra::SVector<f64, 25> while it should return a nalgebra::SMatrix<f64, 25, 25>.

Expected behavior The generated function should use the correct return type.

Environment (please complete the following information):

• OS and version: Pop!_OS 22.04 LTS
• Python version: 3.10
• SymForce Version: 0.10.1

[avsaase]

If I change the function to return a tuple:

def create_matrix() -> tuple[sf.Matrix, sf.Scalar]:
    return sf.Matrix55(
        1, 2, 3, 4, 5,
        6, 7, 8, 9, 10,
        11, 12, 13, 14, 15,
        16, 17, 18, 19, 20,
        21, 22, 23, 24, 25,
    ), sf.Scalar(1)

then the generated function signature has the correct matrix dimensions:

pub fn create_matrix(
        res0: Option<&mut nalgebra::SMatrix<f64, 5, 5>>,
        res1: Option<&mut f64>,
    ) -> ()

I guess it's not directly related to this issue but why does the function not simply return a tuple? And what is the purpose of making the return arguments optional?

[avsaase]

I suspect the problem with the return type in the signature occurs here:

https://github.com/symforce-org/symforce/blob/e270c029f077205a2230fb9266d074dbed3b4f81/symforce/codegen/backends/rust/templates/util/util.jinja#L61-L62

And in the function body here:

https://github.com/symforce-org/symforce/blob/e270c029f077205a2230fb9266d074dbed3b4f81/symforce/codegen/backends/rust/templates/util/util.jinja#L233-L235

I guess these need an if statement like this:

{% if T.SHAPE[0] == 1 or T.SHAPE[1] == 1 %}
    {{ format_vector(T, spec.return_key, spec) }}
{% else %}
    {{ format_matrix(T, spec.return_key, spec) }}
{% endif %}

I'm not able to get the project to build so I can't test this locally. Hopefully someone else more familiar with the project could look in it.

[aaron-skydio]

why does the function not simply return a tuple? And what is the purpose of making the return arguments optional?

Previous discussion here: https://github.com/symforce-org/symforce/pull/405#discussion_r1826870559

[aaron-skydio]

I'm not able to get the project to build so I can't test this locally

Interested at least in what error you're hitting, if it's something we should fix