xlm.modules.encoder

DiffusionTransformerEncoderLayer

Bases: Module

forward(src, src_mask=None, src_key_padding_mask=None)

Parameters:

Name	Type	Description	Default
src		shape (B, T, d_model) for now we will assume T=query_seq_len=key_seq_len	required
src_mask		shape (B*num_heads, T, T) or (T, T). True is attend, False is not attend Note that nn.TransformerEncoderLayer will allow float masks which are added to the attention scores. But we only support boolean masks here.	None
src_key_padding_mask		shape (B, T) or (T). True is attend, False is masked	None

DiffusionTransformerEncoder

Bases: Module

forward(src_tokens, src_lengths)

Parameters:

Name	Type	Description	Default
src_tokens	Tensor	shape (B, T)	required
src_lengths	Optional[Tensor]	shape (B) of type LongTensor	required