For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

matrix_band_part

def matrix_band_part[dtype: DType, int_type: DType, cond_type: DType, rank: Int, simd_width: Int, InputFnType: def[width: Int, rank: Int](IndexList[rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](input_0_fn: InputFnType, input_shape: IndexList[rank], num_lower: TileTensor[int_type, Storage=num_lower.Storage, address_space=num_lower.address_space, linear_idx_type=num_lower.linear_idx_type], num_upper: TileTensor[int_type, Storage=num_upper.Storage, address_space=num_upper.address_space, linear_idx_type=num_upper.linear_idx_type], exclude: TileTensor[cond_type, Storage=exclude.Storage, address_space=exclude.address_space, linear_idx_type=exclude.linear_idx_type], output: TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], ctx: DeviceContext) where (eq InputFnType.dtype, dtype)

Copies a band of input_0_fn into output, zeroing elements outside the band defined by num_lower and num_upper.

Reads the lower and upper diagonal counts and the exclude flag from the supplied tensors, then dispatches to the implementation with the exclude flag specialized at compile time via unswitch.

Parameters:

• ​dtype (DType): Element DType of the input and output tensors.
• ​int_type (DType): DType of the num_lower and num_upper count tensors.
• ​cond_type (DType): DType of the exclude flag tensor.
• ​rank (Int): Number of dimensions in the input tensor (must be at least 2).
• ​simd_width (Int): SIMD width for the elementwise kernel launch.
• ​InputFnType (def[width: Int, rank: Int](IndexList[rank]) -> SIMD[dtype, width] & ImplicitlyCopyable & RegisterPassable): Callable type returning the input element at a given index as a SIMD vector.
• ​target (StringSlice[ImmStaticOrigin]): Target device for the elementwise kernel (defaults to "cpu").

Args:

• ​input_0_fn (InputFnType): Function returning the input element at a given index.
• ​input_shape (IndexList[rank]): Shape of the input tensor.
• ​num_lower (TileTensor[int_type, Storage=num_lower.Storage, address_space=num_lower.address_space, linear_idx_type=num_lower.linear_idx_type]): Scalar tensor giving the number of lower subdiagonals to keep.
• ​num_upper (TileTensor[int_type, Storage=num_upper.Storage, address_space=num_upper.address_space, linear_idx_type=num_upper.linear_idx_type]): Scalar tensor giving the number of upper superdiagonals to keep.
• ​exclude (TileTensor[cond_type, Storage=exclude.Storage, address_space=exclude.address_space, linear_idx_type=exclude.linear_idx_type]): Scalar tensor flag that, when nonzero, inverts the band mask.
• ​output (TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Mutable output tensor receiving the band-part result.
• ​ctx (DeviceContext): Device context used to launch the elementwise kernel.