Function rand

tensorrt_llm/functional.py:1401–1433

This operation adds a fill layer that generates a random (uniform) tensor with the specified shape and data type. Parameters: shape: Tensor The shape of the tensor needed to be generated. low: float The minimum value (inclusive) of the range used for

(shape: Tensor,
         low: float = 0,
         high: float = 1,
         dtype: Union[str, trt.DataType] = 'float32')

def rand(shape: Tensor,
         low: float = 0,
         high: float = 1,
         dtype: Union[str, trt.DataType] = 'float32') -> Tensor:
    '''
    This operation adds a fill layer that generates a random (uniform) tensor with the specified shape and data type.

    Parameters:
        shape: Tensor
            The shape of the tensor needed to be generated.
        low: float
            The minimum value (inclusive) of the range used for random.
        high: float
            The maximum value (inclusive) of the range used for random.
        dtype: Union[str, trt.DataType]
            The desired data type for the output tensor.
    Returns:
        The generated random tensor produced by the fill layer.
    '''
    # NOTE: DISABLED FOR NOW UNTIL THE FILL LAYER (RANDOM_UNIFORM) in TRT IS FIXED
    assert False, "The rand() op is temporarily disabled."
    low = constant(fp32_array(low))
    high = constant(fp32_array(high))
    trt_dtype = dtype if isinstance(dtype,
                                    trt.DataType) else str_dtype_to_trt(dtype)

    layer = default_trtnet().add_fill([0], trt.FillOperation.RANDOM_UNIFORM,
                                      trt_dtype)

    layer.set_input(0, shape.trt_tensor)
    layer.set_input(1, low.trt_tensor)
    layer.set_input(2, high.trt_tensor)
    return _create_tensor(layer.get_output(0), layer)


def categorical_sample(probs: Tensor, rand_data: Tensor = None) -> Tensor: