Mojo function

broadcast_impl

broadcast_impl[rank: Int, output_shape: DimList, input_shape: DimList, type: DType](axis: Int, output: NDBuffer[type, rank, origin, output_shape], input: NDBuffer[type, rank, origin, input_shape], 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 (NDBuffer[type, rank, origin, output_shape]): The output buffer.
input (NDBuffer[type, rank, origin, input_shape]): 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.