Collate function for tokenized audio and text. Merges variable-length audio/text into a single padded tensor.

Forward pass for Sesame's CSM model. This will be added to the model with `model.forward = types.MethodType(forward, model)`

Args: x: shape (b, c, t)

Args: x: shape (b, c, t)

Args: x: shape (b, t, c) x_lens: shape (b,)

Args: x: shape (b, t, c) x_lens: shape (b,)

Forward pass of the ISTFTHead module. Args: x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,

Args: x: shape (b, t, c) x_lens: shape (b,)

Args: attention_mask: Bool mask or float mask. Bool mask, True indicates should attend. Float mask is added to the attention scor

Args: x: shape (b, c, t) cache: shape (b, c_in+c_out, t=2)

ISTFTHead can be adapted in streaming inference without retraining. Args: x: shape (B, T, C) cache: shape (B, N, T=3)

Args: x: input sequence at 50Hz, length should be multiples of 4

Forward self-attention with kv cache. Args: hidden_states: shape (b, t, c) kv_cache: shape (b, nh, t, c*2)

Forward self-attention with kv cache. Args: hidden_states: shape (b, t, c) kv_cache: shape (b, nh, t, c*2)

_summary_ Args: dataset_dir (_type_): _description_

_summary_ Args: logits (torch.Tensor): _description_ n (float): _description_ temperature (float): _description_ Ret