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 of this instruction is specified by the dtype field. The output data type must match the input data type of data if input data type is any of the integer types. Otherwise, output data type can be any valid NKI data types. If output data type is not specified, it is default to be the same as input data type.

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