jax.experimental.pallas.dslice

jax.experimental.pallas.dslice(start, size=None, stride=None)

Constructs a Slice from a start index and a size.

The semantics of dslice mirror those of the builtin slice type:

• dslice(None) is :

• dslice(j) is :j

• dslice(i, j) is i:i+j

• dslice(i, j, stride) is i:i+j:stride

Parameters:

• start (int | Array | None)

• size (int | Array | None)

• stride (int | None)

Return type:

slice | Slice