jax_stack

probly.representation.jax_functions.jax_stack(arrays: Sequence[JaxArrayLike], /, axis: int = 0, out: None = None, dtype: DTypeLike | None = None) → jax.Array[source]

Join a sequence of arrays along a new axis, mirroring jax.numpy.stack.

Parameters:

  • arrays – The arrays to stack.

  • axis – The axis of the result along which the inputs are stacked.

  • out – Unsupported by JAX, must be None.

  • dtype – The desired data type of the result.

Returns:

The stacked array.