jax.lax.zeros_like_array