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

Mojo function

fused_concat

def fused_concat[dtype: DType, rank: Int, input_fn: def[input_index: Int, width: Int, _rank: Int, alignment: Int = Int(1)](IndexList[_rank]) capturing thin -> SIMD[dtype, width], output_0_fn: def[c_type: DType, rank: Int, width: SIMDLength = SIMDLength(1), *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None, output_layout: TensorLayout, *, axis: Int, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu")](input_shapes: StaticTuple[IndexList[rank]], output: TileTensor[dtype, output_layout, Storage=output.Storage], ctx: DeviceContext)

Concatenates inputs produced by input_fn along axis into output, applying output_0_fn to each element.

Instead of reading from concrete input tensors, the fused variant drives the concat from a caller-supplied input_fn that produces values on demand, allowing the concat to be fused with preceding elementwise operations. Dispatches to the CPU or GPU fused implementation based on target and returns early when the output tensor is empty.

Parameters:

dtype (DType): Element type of the input and output tensors.
rank (Int): Number of dimensions in the input and output tensors.
input_fn (def[input_index: Int, width: Int, _rank: Int, alignment: Int = Int(1)](IndexList[_rank]) capturing thin -> SIMD[dtype, width]): Function that produces input element values on demand, indexed by input position.
output_0_fn (def[c_type: DType, rank: Int, width: SIMDLength = SIMDLength(1), *, alignment: Int = Int(1)](IndexList[rank], SIMD[c_type, width]) capturing thin -> None): Epilogue function applied to each produced element before storing.
output_layout (TensorLayout): Layout type of the output tensor.
axis (Int): Axis along which to concatenate the inputs.
target (StringSlice[ImmStaticOrigin]): Target device to dispatch to (defaults to "cpu").

Args:

input_shapes (StaticTuple[IndexList[rank]]): Static tuple of per-input shapes describing each logical input that input_fn produces.
output (TileTensor[dtype, output_layout, Storage=output.Storage]): Destination tensor that receives the concatenated result.
ctx (DeviceContext): Device context used to schedule the work.