Mojo function

broadcast_impl

broadcast_impl[type: DType](axis: Int, output: LayoutTensor[type, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], input: LayoutTensor[type, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], input_prev_axis_stride: Int, output_prev_axis_stride: Int, input_offset: Int, output_offset: Int, rightmost_broadcast_axis: Int)

For each axis of input ∈ [axis, rank), if the dimension is 1, duplicate the data at each index of the corresponding axis in output, otherwise copy over the entire axis to the corresponding axis in output.

Args:

axis (Int): The axis value.
output (LayoutTensor[type, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): The output buffer.
input (LayoutTensor[type, layout, origin, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): The input buffer.
input_prev_axis_stride (Int): The stride at axis axis - 1 for input.
output_prev_axis_stride (Int): The stride at axis axis - 1 for output.
input_offset (Int): The offset at which we start copying data from.
output_offset (Int): The offset at which we start copying data to.
rightmost_broadcast_axis (Int): The largest axis at which we need to duplicate input data.