Table of Contents

Mojo function

mean

fn def mean[dtype: DType, input_fn: def[width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width], output_fn: def[width: SIMDLength, rank: Int](IndexList[rank], SIMD[dtype, width]) capturing thin -> None, /, target: StringSpan[ImmStaticOrigin] = StringSpan("cpu"), *, reduce_dim: Int](input_shape: Coord, output_shape: Coord[input_shape.element_types])

Computes the mean across the input and output shape.

This performs the mean computation on the domain specified by input_shape, loading the inputs using the input_fn. The results' domain is output_shape which are stored using the output_fn.

Parameters:

  • dtype (DType): The dtype of the input and output.
  • input_fn (def[width: Int, rank: Int](IndexList[rank]) capturing thin -> SIMD[dtype, width]): The function to load the input.
  • output_fn (def[width: SIMDLength, rank: Int](IndexList[rank], SIMD[dtype, width]) capturing thin -> None): The function to store the output.
  • target (StringSpan[ImmStaticOrigin]): The target to run on.
  • reduce_dim (Int): The axis to perform the mean on.

Args:

  • input_shape (Coord): The input shape.
  • output_shape (Coord[input_shape.element_types]): The output shape.

Raises:

If the operation fails.

fn def mean[dtype: DType](src: Span[Scalar[dtype]]) -> Scalar[dtype]

Computes the mean value of the elements in a buffer.

Parameters:

  • dtype (DType): The dtype of the input.

Args:

  • src (Span[Scalar[dtype]]): The buffer of elements for which the mean is computed.

Returns:

Scalar[dtype]: The mean value of the elements in the given buffer.

Raises:

If the operation fails.

fn def mean[dtype: DType, input_fn_1d: def[dtype_: DType, width: Int](idx: Int) capturing thin -> SIMD[dtype_, width]](length: Int) -> Scalar[dtype]

Computes the arithmetic mean of values generated by a function.

Parameters:

  • dtype (DType): The data type of the elements.
  • input_fn_1d (def[dtype_: DType, width: Int](idx: Int) capturing thin -> SIMD[dtype_, width]): A function that generates SIMD values at each index.

Args:

  • length (Int): The number of elements to average.

Returns:

Scalar[dtype]: The mean value. For integral types, uses integer division.

Raises:

To comply with how generators are used in this module.

Edit this page