jax.extend.linear_util.merge_linear_aux