Skip to main content
Log in

Mojo function

foreach

foreach[type: DType, rank: Int, //, func: fn[Int](IndexList[$1|1]) capturing -> SIMD[$1|0, $0], synchronous: Bool = False, target: StringLiteral = "cpu", simd_width: Int = get_kernel_simd_width[::DType,::StringLiteral]()](tensor: ManagedTensorSlice[type, rank], ctx: MojoCallContextPtr)

Apply the function func to each element of the tensor slice.

Parameters:

  • type (DType): The data type of the elements in the tensor slice.
  • rank (Int): The rank of the tensor slice.
  • func (fn[Int](IndexList[$1|1]) capturing -> SIMD[$1|0, $0]): The function to apply to each element of the tensor slice.
  • synchronous (Bool): True to run the custom op synchronously in the runtime (defaults to False).
  • target (StringLiteral): A StringLiteral indicating the type of the target device (e.g. "CPU", "CUDA").
  • simd_width (Int): The SIMD width for the target (usually leave this as its default value).

Args:

  • tensor (ManagedTensorSlice[type, rank]): The output tensor slice which receives the return values from func.
  • ctx (MojoCallContextPtr): The call context (forward this from the custom operation).