jax.experimental.pallas.tpu.emit_pipeline

jax.experimental.pallas.tpu.emit_pipeline(body, *, grid, in_specs=(), out_specs=(), tiling=None, core_axis=None, core_axis_name=None, dimension_semantics=None, trace_scopes=True, no_pipelining=False, _explicit_indices=False)

Parameters:

• grid (tuple[int | jax.Array, ...])

• tiling (Tiling | None)

• core_axis (tuple[int, ...] | int | None)

• core_axis_name (tuple[str, ...] | str | None)

• dimension_semantics (tuple[GridDimensionSemantics, ...] | None)

• trace_scopes (bool)

• no_pipelining (bool)

• _explicit_indices (bool)