Mojo function
foreach
foreach[dtype: DType, rank: Int, //, func: def[width: Int](IndexList[rank]) capturing -> SIMD[dtype, width], *, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu"), simd_width: Int = get_kernel_simd_width[dtype, target](), _trace_name: StringSlice[StaticConstantOrigin] = StringSlice("mogg.for_each")](tensor: ManagedTensorSlice[static_spec=tensor.static_spec], ctx: DeviceContext)
Apply the function func to each element of the tensor slice.
Parameters:
- βdtype (
DType): The data type of the elements in the tensor slice. - βrank (
Int): The rank of the tensor slice. - βfunc (
def[width: Int](IndexList[rank]) capturing -> SIMD[dtype, width]): The function to apply to each element of the tensor slice. - βtarget (
StringSlice[StaticConstantOrigin]): Indicates the type of the target device (e.g. "cpu", "gpu"). - βsimd_width (
Int): The SIMD width for the target (usually leave this as its default value). - β_trace_name (
StringSlice[StaticConstantOrigin]): Name of the executed operation displayed in the trace_description.
Args:
- βtensor (
ManagedTensorSlice[static_spec=tensor.static_spec]): The output tensor slice which receives the return values fromfunc. - βctx (
DeviceContext): The call context (forward this from the custom operation).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!