Skip to content

init

Functions:

Name Description
scaled_kaiming_uniform_

Initialize a tensor with a Kaiming distribution scaled by initialization_scale.

supports_flash_attention

Check if a device supports flash attention.

scaled_kaiming_uniform_

scaled_kaiming_uniform_(
    t: Tensor, initialization_scale: float
) -> None

Initialize a tensor with a Kaiming distribution scaled by initialization_scale.

Parameters:

Name Type Description Default

t

Tensor

The tensor to initialize.

required

initialization_scale

float

The amount to scale the initialization by.

required

supports_flash_attention

supports_flash_attention(device: device) -> bool

Check if a device supports flash attention.

Note

Taken from https://github.com/huggingface/transformers/issues/28188#issuecomment-1906901375.

Parameters:

Name Type Description Default

device

device

The device to check, typically a CUDA device.

required

Returns:

Type Description
bool

True if the device supports flash attention, False otherwise.