Mojo function

broadcast_impl

broadcast_impl[dtype: DType](axis: Int, output: TileTensor[dtype, linear_idx_type=output.linear_idx_type, element_size=output.element_size], input: TileTensor[dtype, linear_idx_type=input.linear_idx_type, element_size=input.element_size], 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 (TileTensor[dtype, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): The output buffer.
input (TileTensor[dtype, linear_idx_type=input.linear_idx_type, element_size=input.element_size]): 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.