jax_expand_dims

probly.representation.jax_functions.jax_expand_dims(a: JaxArrayLike, /, axis: int | Sequence[int]) → jax.Array

Insert axes of length one, mirroring jax.numpy.expand_dims.

Parameters:

• a – The array to expand.

• axis – The position or positions of the new axes.

Returns:

The expanded array.