jax.lax.expand\_dims ==================== .. currentmodule:: jax.lax .. autofunction:: expand_dims