Codegen not using correct object dimensions

[gftabor]

Describe the bug Python codegen with matrix doesn't expect correct dimensions? The input dimensions were correctly put in args 10x4 for A and B but the codegen calls element 30. When I pass a matrix of 10x4 to my generated function it throws an error because 30 isn't possible, it seems the codegen expects me to flatten the array before passing it in.

` def error_factor(A, B, x):

type: (numpy.ndarray, numpy.ndarray, T.Sequence[float]) -> T.Tuple[T.List[float], numpy.ndarray, numpy.ndarray, T.List[float]]

"""
This function was autogenerated from a symbolic function. Do not modify by hand.

Symbolic function: error

Args:
    A: Matrix10_4
    B: Matrix10_4
    x: Matrix41

Outputs:
    res: Matrix11
    jacobian: (1x4) jacobian of res wrt arg x (4)
    hessian: (4x4) Gauss-Newton hessian for arg x (4)
    rhs: (4x1) Gauss-Newton rhs for arg x (4)
"""

# Total ops: 2559

# Input arrays

# Intermediate terms (515)
_tmp0 = A[0] * B[0] + A[10] * B[10] + A[20] * B[20] + A[30] * B[30]
_tmp1 = A[0] * _tmp0`

To Reproduce

A = np.random.random((size,4))
B = np.random.random((size,4))

x = np.random.random((4,1))

x_s = sf.Matrix.zeros(4, 1).symbolic("x")
x_s2 = symforce.geo.matrix.Matrix41(x)

def error(A: sf.Matrix(size,4),B: sf.Matrix(size,4),x: sf.Matrix(4,1)) -> sf.V1:
    stuff = A * x +  A * B.T * A * x
    return sf.V1(stuff.T * stuff)

s = time.time()
from symforce.codegen import Codegen, CppConfig, PythonConfig

codegen = Codegen.function(error, config=PythonConfig())

codegen_linearization = codegen.with_linearization(
    which_args=["x"]
)

metadata = codegen_linearization.generate_function()

Expected behavior Codegen should index elements correctly from the input shape.

Environment (please complete the following information):

• OS and version: Ubuntu Linux 20.04
• Python version 3.9
• SymForce Version 0.5.0

[gftabor]

Just updating this, the indexing is correct with C++ codegen which I imagine is why it was never noticed.

void ErrorFactor(const Eigen::Matrix<Scalar, 10, 4>& A, const Eigen::Matrix<Scalar, 10, 4>& B,
                 const Eigen::Matrix<Scalar, 4, 1>& x,
                 Eigen::Matrix<Scalar, 1, 1>* const res = nullptr,
                 Eigen::Matrix<Scalar, 1, 4>* const jacobian = nullptr,
                 Eigen::Matrix<Scalar, 4, 4>* const hessian = nullptr,
                 Eigen::Matrix<Scalar, 4, 1>* const rhs = nullptr) {
  // Total ops: 2559

  // Input arrays

  // Intermediate terms (515)
  const Scalar _tmp0 =
      A(0, 0) * B(0, 0) + A(0, 1) * B(0, 1) + A(0, 2) * B(0, 2) + A(0, 3) * B(0, 3);
  const Scalar _tmp1 = A(0, 0) * _tmp0;