def METHOD_NAME(c, xs, init_mode=False):
  # inject constants
  xs = jax.tree_util.tree_map(
      lambda ax, arg, x: (arg if ax is broadcast else x), in_axes, args, xs
  )
  broadcast_out, c, ys = fn(broadcast_in, c, *xs)
  if init_mode:
    ys = jax.tree_util.tree_map(
        lambda ax, y: (y if ax is broadcast else ()), out_axes, ys
    )
    return broadcast_out, ys
  else:
    ys = jax.tree_util.tree_map(
        lambda ax, y: (() if ax is broadcast else y), out_axes, ys
    )
    return c, ys