nki.isa.dropout#

nki.isa.dropout(dst, data, prob, name=None)[source]#

Randomly replace some elements of the input tile data with zeros based on input probabilities using Vector Engine. The probability of replacing input elements with zeros (i.e., drop probability) is specified using the prob field: - If the probability is 1.0, all elements are replaced with zeros. - If the probability is 0.0, all elements are kept with their original values.

The prob field can be a scalar constant or a tile of shape (data.shape[0], 1), where each partition contains one drop probability value. The drop probability value in each partition is applicable to the input data elements from the same partition only.

Data type of the input data tile can be any valid NKI data types (see Supported Data Types for more information). However, data type of prob has restrictions based on the data type of data:

  • If data type of data is any of the integer types (e.g., int32, int16), prob data type must be float32

  • If data type of data is any of the float types (e.g., float32, bfloat16), prob data can be any valid float type

The output data type dst.dtype must match the input data type data.dtype.

Parameters:

  • dst – an output tile of the dropout result

  • data – the input tile

  • prob – a scalar or a tile of shape (data.shape[0], 1) to indicate the probability of replacing elements with zeros