jax.lax.expand_dims

jax.lax.expand_dims(array, dimensions)

Insert any number of size 1 dimensions into an array.

Parameters:

• array (Union[Array, ndarray, bool_, number, bool, int, float, complex])

• dimensions (Sequence[int])

Return type:

Array