This document is relevant for: Trn2, Trn3
nki.language.mean#
- nki.language.mean(x, axis, dtype=None, keepdims=False)[source]#
Arithmetic mean along the specified axis (or axes) of the input.
((Similar to numpy.mean))
Warning
This API is experimental and may change in future releases.
- Parameters:
x – a tile.
axis – int or tuple/list of ints. Free dimensions can only be the last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. The partition dimension (0) is also supported on its own, reducing across partitions; it cannot be mixed with free dimensions in one call.
dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.
keepdims – if True, the reduced axes are kept as size-one dimensions. It has no effect for
axis=0: dimension 0 is structurally the partition dimension of a tile, so it is always retained as size one rather than dropped.
- Returns:
a tile with the average of elements along the provided axis. Float32 intermediate values are used for the computation.
This document is relevant for: Trn2, Trn3