IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

elementwise

def elementwise[func: def[width: Int, alignment: Int = Int(1)](Coord[**?]) capturing thin -> None, simd_width: Int, *, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), _trace_description: StringSlice[ImmStaticOrigin] = StringSlice("elementwise")](shape: Int, context: DeviceContext)

Executes func[width, rank](indices), possibly as sub-tasks, for a suitable combination of width and indices so as to cover shape. Returns when all sub-tasks have completed.

Parameters:

  • ​func (def[width: Int, alignment: Int = Int(1)](Coord[**?]) capturing thin -> None): The body function.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[ImmStaticOrigin]): The target to run on.
  • ​_trace_description (StringSlice[ImmStaticOrigin]): Description of the trace.

Args:

  • ​shape (Int): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.

def elementwise[func: def[width: Int, alignment: Int = Int(1)](Coord[**?]) capturing thin -> None, simd_width: Int, *, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), _trace_description: StringSlice[ImmStaticOrigin] = StringSlice("elementwise")](shape: Coord, context: DeviceContext)

Executes func[width, rank](indices), possibly as sub-tasks, for a suitable combination of width and indices so as to cover shape. Returns when all sub-tasks have completed.

Parameters:

  • ​func (def[width: Int, alignment: Int = Int(1)](Coord[**?]) capturing thin -> None): The body function.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[ImmStaticOrigin]): The target to run on.
  • ​_trace_description (StringSlice[ImmStaticOrigin]): Description of the trace.

Args:

  • ​shape (Coord): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.

def elementwise[FuncType: def[width: Int, alignment: Int = Int(1)](Coord[**?]) -> None & ImplicitlyCopyable & RegisterPassable, //, simd_width: Int, *, target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), _trace_description: StringSlice[ImmStaticOrigin] = StringSlice("elementwise")](func: FuncType, shape: Coord, context: DeviceContext)

Unified-closure entry point for elementwise (DeviceContext).

Accepts a parametric body (already in unified-closure form, with explicit captures) and dispatches to _elementwise_impl. rank and FuncType are inferred from the runtime shape and func arguments, so simd_width is the only explicit positional parameter β€” callers can write elementwise[N](func, shape, ctx).

Parameters:

  • ​FuncType (def[width: Int, alignment: Int = Int(1)](Coord[**?]) -> None & ImplicitlyCopyable & RegisterPassable): A parametric callable taking IndexList[rank] and template parameters width, rank, alignment.
  • ​simd_width (Int): The SIMD vector width to use.
  • ​target (StringSlice[ImmStaticOrigin]): The target to run on.
  • ​_trace_description (StringSlice[ImmStaticOrigin]): Description of the trace.

Args:

  • ​func (FuncType): The body closure value.
  • ​shape (Coord): The shape of the buffer.
  • ​context (DeviceContext): The device context to use.

Raises:

If the operation fails.