Converting code from Ceres to Jax, dealing with array of structures

[isaacgerg]

In Ceres, residuals are often represented as part of a class. In python this looks like:

class Residual:
   def __init__(self, p1, p2, p3)
      self._p1 = p1, etc
   def residual(self, trial_value):
        return equation_of_trial_value_and_self.p1_self._p2_etc

So the residuals are essentially array of structures.

In Jax, to process these, a natural way would be something like:

def function_to_minimize(trial_value, residuals):
   error = 0
   for k in num_residuals:
       error = error + residuals[k].residual(trial_value)
   return error

However several issues arise from this:

1. Residual is not a numpy array so residuals cannot be compiled via jax.jit(function_to_minimize), # Argument '<main.Residual object at 0x7f8ab63aac18>' of type <class 'main.Residual'> is not a valid JAX type
2. residuals is a list of Residual and cannot be accessed via .scan

I am asking if there is a pattern that I can use which closely matches the canonical formulation used by Ceres. It is true that I can rewrite everything but if I can keep the formulation, it would greatly reduce the changes of error.

Thanks!

[shoyer]

Making Residual a custom pytree would probably fix your first issue: https://jax.readthedocs.io/en/latest/pytrees.html

I don't think there's an easy fix for (2). scan fundamentally requires an array with a leading axis to scan over. On the other hand, if using a list gives you good enough performance, you can still use normal for loops in JAX as well (this often results in slow compilation times if you use jit, though)

[isaacgerg]

Thanks!

Interestingly enough, scan does not work but the code works fine if i use the psuedo-code the documentation gives for scan. See below.

Here is my pytree

jax.tree_util.register_pytree_node(
    Residual,
    lambda xs: (((xs.tx, xs.rx1, xs.rx2, xs.vx, xs.p1_attitude, xs.p2_attitude, xs.altitude, xs.delay, xs.dt, xs.t_between_pings, xs.c, xs.weight)), None),  # tell JAX how to unpack to an iterable
    lambda _, xs: Residual(*xs)       # tell JAX how to pack back into a Residual
)

I pass to my error function a list of Residual objects. The following works fine:

        def scan(f, init, xs, length=None):  # This definition from the Jax documentation.
            if xs is None:
                xs = [None] * length
            carry = init
            ys = []
            for x in xs:
                carry, y = f(carry, x)
                ys.append(y)
            return carry, jnp.stack(ys)        
        carry, ys = scan(fun, carry, residuals, jnp.array(num_measurements))

but if i do
carry, ys = jax.lax.scan(fun, carry, residuals)

It gives me out of range since inside of scan, it flattens the whole array giving a length of 924 instead of 77. (i.e. 77 params * 12 args for the input of Residual class = 924 params).

So scan doesn't quite work the way the documentation says it does. Unless I am doing something wrong with my tree definition.? The prototype for the constructor of residual is

class Residual(object):
    def __init__(self, tx, rx1, rx2, vx, p1_attitude, p2_attitude, altitude, delay, dt, t_between_pings, c, weight):

[isaacgerg]

I would expect that .scan would call my packing function in the pytree definition to make my Residual object again before passing body of the loop.

        def scan(f, init, xs, length=None):  
            if xs is None:
                xs = [None] * length
            carry = init
            ys = []
            for x in xs:  # I would think that xs is a list or array of Residual's, not an array of parameters input to Residual objects
                carry, y = f(carry, x)
                ys.append(y)
            return carry, jnp.stack(ys)  

[isaacgerg]

@shoyer Is .scan supposed to return an unflattened version of the objects prior to running? It seems like this shouldn't be difficult to do.

[isaacgerg]

Answered by #3826