jax.experimental.pallas.Slice

class jax.experimental.pallas.Slice(start, size, stride=1)

A slice with a start index and a size.

Both start index and size can either be static, i.e. known at tracing and compilation time, or dynamic.

Parameters:

• start (int | Array)

• size (int | Array)

• stride (int)

__init__(start, size, stride=1)

Parameters:

• start (int | Array)

• size (int | Array)

• stride (int)

Return type:

None

Methods

__init__(start, size[, stride])
from_slice(slc, size)
tree_flatten()
tree_unflatten(aux_data, children)

Attributes

is_dynamic_size
is_dynamic_start
stride
start
size