Constants in scan().

[MRiabov]

Sup,
I recently had to write a module where we had 25 scans in a single file... Well, it was not so simple because we had 5 variables in carry - and most of them were constants. Minimal reproducible example here:

def main_func(a_const, b_const, c_const, d_const, e_const, variable, array:jnp.ndarray)
     def compute_variables(carry, x):
          a_const, b_const, c_const, d_const, e_const, variable= carry # it was even longer and more difficult to understand since variables are very different in function and naming 
          array=x
          # add a dozen lines of logic here
          variable = jnp.concat(x+a_const-b_const*variable, variable)
          return a_const, b_const, c_const, d_const, e_const, variable # again, over 25 scans this was sickening and more difficult to read.
          
    scan(compute_variable, init=(a_const, b_const, c_const, d_const, e_const, variable), xs=array)

I believe in ML there are also hyperparameters which need to be passed and needelessly optimized for by jit in every line. If there is a constants, XLA may not bother optimizing functions for them in the same way as static_argnums makes functions faster..

Would scan maybe get an optional consts parameter in which the constant parameters would be passed?
It's a QOL feature which also may speed up scan and other loops computation a little bit.

Is it big enough to be implemeneted?
(same goes for fori_loop and while_loops)

[jakevdp]

You could pass the constants by closure instead; modifying your example, it might look something like this:

def main_func(a_const, b_const, c_const, d_const, e_const, variable, array:jnp.ndarray)
    def compute_variables(variable, x):
        # add a dozen lines of logic here
        return (x+a_const-b_const*variable, variable)
    scan(compute_variable, init=variable), xs=array)

The general approach is: you can reference a_const and b_const from the outer scope without explicitly passing via the carry. Is that more-or-less what you want?