Mojo function

sliced_add

sliced_add[dtype: DType, //, target: StringSlice[StaticConstantOrigin]](c: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], a: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], b: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], lora_end_idx: TileTensor[DType.int64, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], ctx: Optional[DeviceContext])

Adds tensors a and b element-wise for rows < lora_end_idx, otherwise copies a.

This is used for LoRA where only some sequences have LoRA applied. For rows in [0, lora_end_idx): c = a + b For rows in [lora_end_idx, batch_seq_len): c = a

Args:

c (TileTensor): Output tensor.
a (TileTensor): First input tensor.
b (TileTensor): Second input tensor.
lora_end_idx (TileTensor): Scalar tensor with end index of LoRA token portion (rows to apply add).
ctx (Optional): Device context for GPU operations.