so-vits-svc

Running

App Files Files Community

so-vits-svc / vencoder /dphubert /components.py

aoxiang1221

update

85ce65e 9 months ago

raw history blame

No virus

58.6 kB

	"""Building blocks for speech SSL models supporting pruning.

	Originally from:
	https://github.com/pytorch/audio/blob/main/torchaudio/models/wav2vec2/components.py

	"""

	import math
	from collections import defaultdict
	from typing import List, Optional, Tuple

	import torch
	from torch import Tensor, nn
	from torch.nn import Module

	from .hardconcrete import HardConcrete
	from .pruning_utils import (
	prune_conv1d_layer,
	prune_layer_norm,
	prune_linear_layer,
	)


	def _init_transformer_params(module):
	"""
	Initialize the weights of Transformer module in Wav2Vec2/HuBERT.

	If the module is ``nn.Linear``, normalize the weight with mean 0 and standard deviation 0.02.
	If ``bias`` is set to ``True`` in the module, set ``bias`` to 0.

	If the module is ``nn.Embedding``, normalize the weight with mean 0 and standard deviation 0.02.
	If ``padding_idx`` is not None, set the weight of padding to 0.

	Note:
	Ths method corresponds to
	`init_bert_params
	<https://github.com/facebookresearch/fairseq/blob/main/fairseq/modules/transformer_sentence_encoder.py#L21>`__
	in the original ``fairseq`` implementation.
	"""

	def normal_(data):
	data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))

	if isinstance(module, nn.Linear):
	normal_(module.weight.data)
	if module.bias is not None:
	module.bias.data.zero_()
	if isinstance(module, nn.Embedding):
	normal_(module.weight.data)
	if module.padding_idx is not None:
	module.weight.data[module.padding_idx].zero_()


	class LayerNorm(nn.LayerNorm):
	"""Layer norm with transpose"""

	def forward(self, input: Tensor) -> Tensor:
	x = input.transpose(-2, -1)
	x = nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
	x = x.transpose(-2, -1)
	return x


	class ConvLayerBlock(Module):
	"""Convolution unit of FeatureExtractor"""

	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	kernel_size: int,
	stride: int,
	bias: bool,
	layer_norm: Optional[Module],
	prune_conv_channels: bool = False,
	):
	super().__init__()
	self.kernel_size = kernel_size
	self.stride = stride
	self.layer_norm = layer_norm
	self.conv = nn.Conv1d(
	in_channels=in_channels,
	out_channels=out_channels,
	kernel_size=kernel_size,
	stride=stride,
	bias=bias,
	)

	if prune_conv_channels:
	self.hard_concrete = HardConcrete(n_in=out_channels, init_mean=0.01)
	else:
	self.hard_concrete = None

	def forward(
	self,
	x: Tensor,
	length: Optional[Tensor],
	) -> Tuple[Tensor, Optional[Tensor]]:
	"""
	Args:
	x (Tensor): Shape: ``[batch, in_channels, in_frame]``.
	length (Tensor or None, optional): Shape ``[batch, ]``.
	Returns:
	Tensor: Shape ``[batch, out_channels, out_frames]``.
	Optional[Tensor]: Shape ``[batch, ]``.
	"""
	x = self.conv(x)
	if self.layer_norm is not None:
	x = self.layer_norm(x)
	x = nn.functional.gelu(x)

	if self.hard_concrete is not None:
	channel_mask = self.hard_concrete() # hard concrete mask, (out_channels,)
	x = x * channel_mask.unsqueeze(-1)

	if length is not None:
	length = torch.div(length - self.kernel_size, self.stride, rounding_mode="floor") + 1
	# When input length is 0, the resulting length can be negative. So fix it here.
	length = torch.max(torch.zeros_like(length), length)
	return x, length

	def get_num_params_and_out_channels(self, in_channels):
	if self.hard_concrete is not None:
	out_channels = self.hard_concrete.l0_norm()
	else:
	out_channels = self.conv.out_channels

	num_params = in_channels * out_channels * self.kernel_size
	if self.conv.bias is not None:
	num_params += out_channels
	if self.layer_norm is not None:
	num_params += out_channels * 2

	return num_params, out_channels


	class FeatureExtractor(Module):
	"""Extract features from audio

	Args:
	conv_layers (nn.ModuleList):
	convolution layers
	"""

	def __init__(
	self,
	conv_layers: nn.ModuleList,
	):
	super().__init__()
	self.conv_layers = conv_layers

	# NOTE: a dummy weight used to save the soft mask of the last conv layer
	self.dummy_weight = nn.Parameter(
	torch.ones(conv_layers[-1].conv.out_channels, dtype=torch.float32),
	requires_grad=False
	)

	def forward(
	self,
	x: Tensor,
	length: Optional[Tensor],
	) -> Tuple[Tensor, Optional[Tensor]]:
	"""
	Args:
	x (Tensor):
	Input Tensor representing a batch of audio,
	shape: ``[batch, time]``.
	length (Tensor or None, optional):
	Valid length of each input sample. shape: ``[batch, ]``.

	Returns:
	Tensor:
	The resulting feature, shape: ``[batch, frame, feature]``
	Optional[Tensor]:
	Valid length of each output sample. shape: ``[batch, ]``.
	"""
	if x.ndim != 2:
	raise ValueError("Expected the input Tensor to be 2D (batch, time), " "but received {list(x.shape)}")

	x = x.unsqueeze(1) # (batch, channel==1, frame)
	for layer in self.conv_layers:
	x, length = layer(x, length) # (batch, feature, frame)
	x = x.transpose(1, 2) # (batch, frame, feature)
	x = x * self.dummy_weight
	return x, length

	def get_num_params_and_final_out_channels(self):
	in_channels = 1
	num_params = 0
	for layer in self.conv_layers:
	layer_params, in_channels = layer.get_num_params_and_out_channels(in_channels)
	num_params += layer_params

	num_params += in_channels # dummy weight

	return num_params, in_channels

	def prune(self):
	""""Prune conv layers and dummy weight based on hardconcrete parameters.
	This is an in-place operation.
	"""
	new_config = [] # [(output_channel, kernel_size, stride), ...]
	for idx, layer in enumerate(self.conv_layers):
	if layer.hard_concrete is not None:
	assert not layer.hard_concrete.training
	mask = layer.hard_concrete() # (out_features,)
	index = mask.nonzero().squeeze(-1) # 2D -> 1D
	assert len(index) > 0, f"Conv channels pruned to zero at index {idx}"
	new_config.append(
	(len(index), layer.kernel_size, layer.stride)
	)

	# prune the current layer
	prune_conv1d_layer(layer.conv, index, "output")
	if layer.layer_norm is not None:
	prune_layer_norm(layer.layer_norm, index)

	# prune the next layer
	if idx == len(self.conv_layers) - 1:
	self.dummy_weight.data *= mask
	self.dummy_weight = nn.Parameter(
	self.dummy_weight.index_select(0, index).clone().detach(), requires_grad=False
	)
	else:
	self.conv_layers[idx+1].conv.weight.data *= mask.unsqueeze(-1)
	prune_conv1d_layer(self.conv_layers[idx+1].conv, index, dim="input")

	layer.hard_concrete = None
	else:
	new_config.append(
	(layer.conv.out_channels, layer.kernel_size, layer.stride)
	)
	index = torch.arange(layer.conv.out_channels, dtype=torch.long)

	return new_config, index


	class FeatureProjection(Module):
	"""Layer that connects FeatureExtractor and Encoder

	Projects features to encoder dimension.

	Args:
	in_features (int): Input feature dim.
	out_features (int): Output feature dim.
	dropout (float): Dropout probability.
	"""

	def __init__(
	self,
	in_features: int,
	out_features: int,
	dropout: float,
	):
	super().__init__()
	self.layer_norm = nn.LayerNorm(in_features)
	self.projection = nn.Linear(
	in_features,
	out_features,
	)
	self.dropout = nn.Dropout(dropout)

	def forward(self, x):
	"""
	Args:
	x (Tensor):
	Feature Tensor. shape: ``[batch, frame, in_feature]``
	Returns:
	Tensor: Projected features. ``[batch, frame, out_feature]``.
	"""
	x = self.layer_norm(x)
	x = self.projection(x)
	x = self.dropout(x)
	return x

	def get_num_params(self, in_features):
	return in_features * 2 + (in_features + 1) * self.projection.out_features


	class ConvolutionalPositionalEmbedding(Module):
	"""Positional embedding which is placed at the beginning of Transformer.

	Args:
	embed_dim (int): Feature dimension of the input Tensor.
	kernel_size (int): The number of frames to be use.
	groups (int): The number of groups in feature dimensions.
	"""

	def __init__(
	self,
	embed_dim: int,
	kernel_size: int,
	groups: int,
	):
	super().__init__()
	self.embed_dim = embed_dim
	self.kernel_size = kernel_size
	self.conv = nn.Conv1d(
	in_channels=embed_dim,
	out_channels=embed_dim,
	kernel_size=kernel_size,
	padding=kernel_size // 2,
	groups=groups,
	)

	self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
	self.num_remove: int = 1 if kernel_size % 2 == 0 else 0

	def __prepare_scriptable__(self):
	for hook in self.conv._forward_pre_hooks.values():
	# The hook we want to remove is an instance of WeightNorm class, so
	# normally we would do `if isinstance(...)` but this class is not accessible
	# because of shadowing, so we check the module name directly.
	# https://github.com/pytorch/pytorch/blob/be0ca00c5ce260eb5bcec3237357f7a30cc08983/torch/nn/utils/__init__.py#L3
	if hook.__module__ == "torch.nn.utils.weight_norm" and hook.__class__.__name__ == "WeightNorm":
	torch.nn.utils.remove_weight_norm(self.conv)
	return self

	def forward(self, x):
	"""
	Args:
	x (Tensor): shape ``[batch, frame, feature]``.

	Returns:
	Tensor: The resulting feature. Shape ``[batch, frame, feature]``.
	"""
	x = x.transpose(-2, -1)
	x = self.conv(x)
	if self.num_remove > 0:
	x = x[..., : -self.num_remove]
	x = torch.nn.functional.gelu(x)
	x = x.transpose(-2, -1)
	return x


	class SelfAttention(Module):
	"""Multihead Self Attention module

	Args:
	embed_dim (int): Total dimension of the model.
	num_heads (int): The number of heads.
	dropout (float, optional):
	Dropout probability on attn_output_weights. Default: ``0.0``
	"""

	def __init__(
	self,
	embed_dim: int,
	num_heads: int,
	head_dim: int,
	dropout: float = 0.0,
	prune_heads: bool = False, # whether to prune attention heads
	prune_layer: bool = False, # whether to prune entire attention layers
	):
	super().__init__()

	self.embed_dim = embed_dim
	self.num_heads = num_heads
	self.head_dim = head_dim
	self.dropout = torch.nn.Dropout(dropout)

	self.scaling = self.head_dim**-0.5

	self.k_proj = nn.Linear(embed_dim, num_heads * head_dim, bias=True)
	self.v_proj = nn.Linear(embed_dim, num_heads * head_dim, bias=True)
	self.q_proj = nn.Linear(embed_dim, num_heads * head_dim, bias=True)
	self.out_proj = nn.Linear(num_heads * head_dim, embed_dim, bias=True)

	if prune_heads:
	self.hard_concrete_for_heads = HardConcrete(n_in=num_heads, init_mean=0.01)
	else:
	self.hard_concrete_for_heads = None

	if prune_layer:
	self.hard_concrete_for_layer = HardConcrete(n_in=1, init_mean=0.01)
	else:
	self.hard_concrete_for_layer = None

	def forward(
	self,
	x: Tensor,
	attention_mask: Optional[Tensor] = None,
	position_bias: Optional[Tensor] = None,
	key_padding_mask: Optional[Tensor] = None,
	) -> Tuple[Tensor, Optional[Tensor]]:
	"""
	Args:
	x (Tensor): shape: ``[batch_size, sequence_length, embed_dim]``.
	attention_mask (Tensor or ``None``, optional):
	shape: ``[batch_size, 1, sequence_length, sequence_length]``
	position_bias: Not used. Only for the compatibility with :py:class:`WavLMSelfAttention`.
	key_padding_mask (Tensor or ``None``): Not used. Only for the compatibility with
	:py:class:`WavLMSelfAttention`.
	Returns:
	(Tensor, ``None``): The resulting attention output and ``None`` (necessary for compatibility
	with :py:class:`WavLMSelAttention`).
	Attention output shape: ``[batch, sequence_length, embed_dim]``.
	"""
	if x.ndim != 3 or x.shape[2] != self.embed_dim:
	raise ValueError(
	f"The expected input shape is (batch, sequence, embed_dim=={self.embed_dim}). " f"Found {x.shape}."
	)
	batch_size, length, embed_dim = x.size()

	shape = (batch_size, length, self.num_heads, self.head_dim)
	q = self.q_proj(x).view(*shape).transpose(2, 1) # B, nH, L, Hd
	k = self.k_proj(x).view(*shape).permute(0, 2, 3, 1) # B, nH, Hd, L
	v = self.v_proj(x).view(*shape).transpose(2, 1) # B, nH, L, Hd

	# scale down q to avoid value overflow.
	weights = (self.scaling * q) @ k # B, nH, L, L
	if attention_mask is not None:
	weights += attention_mask
	# subtracting a constant value from the tensor won't change the output of softmax.
	# apply the subtraction to avoid value overflow in torch.nn.functional.softmax.
	# for more details, please see Equation 7 in https://arxiv.org/abs/2112.08778
	weights = weights - weights.max(dim=-1, keepdim=True)[0]

	weights = torch.nn.functional.softmax(weights, dim=-1)
	weights = self.dropout(weights)

	output = weights @ v # B, nH, L, Hd

	if self.hard_concrete_for_heads is not None:
	head_mask = self.hard_concrete_for_heads() # (nH,)
	output = output * head_mask.unsqueeze(-1).unsqueeze(-1)

	output = output.transpose(2, 1).reshape(batch_size, length, self.num_heads * self.head_dim)

	output = self.out_proj(output)

	if self.hard_concrete_for_layer is not None:
	layer_mask = self.hard_concrete_for_layer() # (1,)
	output = output * layer_mask

	return output, None # Necessary for compatibility with WavLMSelAttention

	def get_num_params(self):
	if self.hard_concrete_for_heads is not None:
	num_heads = self.hard_concrete_for_heads.l0_norm()
	else:
	num_heads = self.num_heads
	num_params = (self.embed_dim + 1) * num_heads * self.head_dim * 3 \
	+ (num_heads * self.head_dim + 1) * self.embed_dim

	if self.hard_concrete_for_layer is not None:
	num_params *= self.hard_concrete_for_layer.l0_norm()

	return num_params

	def prune(self):
	new_config = {
	"use_attention": True,
	"num_heads": self.num_heads,
	}
	if self.hard_concrete_for_layer is not None:
	assert not self.hard_concrete_for_layer.training
	layer_mask = self.hard_concrete_for_layer() # (1,)
	self.out_proj.weight.data *= layer_mask
	self.out_proj.bias.data *= layer_mask
	if layer_mask == 0:
	new_config["use_attention"] = False
	self.hard_concrete_for_layer = None

	if self.hard_concrete_for_heads is not None:
	assert not self.hard_concrete_for_heads.training
	head_mask = self.hard_concrete_for_heads() # (num_heads,)
	new_config["num_heads"] = len(head_mask.nonzero())
	if new_config["num_heads"] == 0:
	new_config["use_attention"] = False
	else:
	full_mask = head_mask.repeat_interleave(self.head_dim)
	full_index = full_mask.nonzero().squeeze(-1) # 1D

	prune_linear_layer(self.k_proj, full_index, "output")
	prune_linear_layer(self.v_proj, full_index, "output")
	prune_linear_layer(self.q_proj, full_index, "output")

	self.out_proj.weight.data *= full_mask
	prune_linear_layer(self.out_proj, full_index, "input")
	self.hard_concrete_for_heads = None

	return new_config


	class WavLMSelfAttention(SelfAttention):
	"""Multi-headed self-attention for WavLM model :cite:`chen2022wavlm`.

	Args:
	embed_dim (int): Total dimension of the model.
	num_heads (int): The number of heads.
	dropout (float, optional): Dropout probability on attn_output_weights. (Default: to ``0.0``)
	bias (bool, optional): If ``True``, add bias to input / output projection layers. (Default: ``True``)
	has_relative_attention_bias (bool, optional): If ``True``, apply relative position embedding.
	Necessary in the first encoder layer, but not in the subsequent ones. (Default: ``False``)
	num_buckets (int, optional): Number of buckets for relative position embedding. (Default: ``32``)
	max_distance (int, optional): Naximum distance for relative position embedding. (Default: ``128``)
	gru_rel_pos (bool, optional): If ``True``, apply gated relative position embedding. (Default: ``False``)
	"""

	def __init__(
	self,
	embed_dim: int,
	total_num_heads: int,
	remaining_heads: Optional[List[int]] = None,
	dropout: float = 0.0,
	bias: bool = True,
	has_relative_attention_bias: bool = False,
	num_buckets: int = 32,
	max_distance: int = 128,
	gru_rel_pos: bool = True,
	prune_heads: bool = False,
	prune_layer: bool = False,
	):
	self.total_num_heads = total_num_heads
	if remaining_heads is None:
	self.remaining_heads = list(range(total_num_heads))
	else:
	self.remaining_heads = remaining_heads # list of indices

	self.head_dim = embed_dim // total_num_heads

	super().__init__(embed_dim, len(self.remaining_heads), self.head_dim, dropout, prune_heads, prune_layer)

	self.has_relative_attention_bias = has_relative_attention_bias
	self.num_buckets = num_buckets
	self.max_distance = max_distance

	if has_relative_attention_bias:
	self.rel_attn_embed = nn.Embedding(num_buckets, total_num_heads)
	else:
	self.rel_attn_embed = None

	# override linear layers to customize bias
	self.k_proj = nn.Linear(embed_dim, len(self.remaining_heads) * self.head_dim, bias=bias)
	self.v_proj = nn.Linear(embed_dim, len(self.remaining_heads) * self.head_dim, bias=bias)
	self.q_proj = nn.Linear(embed_dim, len(self.remaining_heads) * self.head_dim, bias=bias)
	self.out_proj = nn.Linear(len(self.remaining_heads) * self.head_dim, embed_dim, bias=bias)

	self.gru_rel_pos = gru_rel_pos
	if self.gru_rel_pos:
	self.gru_rel_pos_linear = nn.Linear(self.head_dim, 8)
	self.gru_rel_pos_const = nn.Parameter(torch.ones(1, total_num_heads, 1, 1))
	self.has_position_bias = True

	def compute_bias(self, query_length: int, key_length: int) -> Tensor:
	"""Compute relative position embeddings for WavLM model.
	Args:
	query_length (int): Query position can take values between 0 and ``query_length - 1``.
	key_length (int): Key position can take values between 0 and ``key_length - 1``.
	Returns:
	Tensor of shape `(num_heads, query_length, key_length)`, relative positions embeddings
	"""
	context_position = torch.arange(query_length, dtype=torch.long)[:, None]
	memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
	relative_position = memory_position - context_position # Shape (query_length, key_length)
	relative_position_bucket = self._relative_positions_bucket(relative_position, bidirectional=True)
	relative_position_bucket = relative_position_bucket.to(self.rel_attn_embed.weight.device)
	values = self.rel_attn_embed(relative_position_bucket) # Shape (query_length, key_length, num_heads)
	values = values.permute([2, 0, 1])
	return values

	def _relative_positions_bucket(self, relative_positions: Tensor, bidirectional: bool = True):
	"""Compute relative position buckets for WavLM model. Computation similar to formula (5) in WavLM
	paper :cite:`chen2022wavlm`.
	Args:
	relative_positions (Tensor): Relative offsets between query and key positions,
	of shape ``(query_length, key_length)``.
	bidirectional (bool): If ``True``, values will be filled both above and below the diagonal in the resulting
	matrix. If ``False``, the elements above the diagonal (i.e. with negative relative offsets) will be set
	to zero. (Default ``True``)
	Returns:
	Tensor of shape ``(query_length, key_length)`` filled bucketed values of with relative positions.
	"""
	num_buckets = self.num_buckets
	max_distance = self.max_distance
	# Shape (query_length, key_length)
	relative_buckets = torch.zeros_like(relative_positions, dtype=torch.long)

	if bidirectional:
	num_buckets = num_buckets // 2
	relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
	relative_positions = torch.abs(relative_positions)
	else:
	relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions))

	max_exact = num_buckets // 2
	is_small = relative_positions < max_exact

	relative_postion_if_large = max_exact + (
	torch.log(relative_positions.float() / max_exact)
	/ math.log(max_distance / max_exact)
	* (num_buckets - max_exact)
	).to(torch.long)
	relative_postion_if_large = torch.min(
	relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1)
	)

	relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large)
	return relative_buckets

	def forward(
	self,
	query: Tensor,
	attention_mask: Optional[Tensor] = None,
	position_bias: Optional[Tensor] = None,
	key_padding_mask: Optional[Tensor] = None,
	) -> Tuple[Tensor, Optional[Tensor]]:
	"""
	Args:
	query (Tensor): Input of shape ``(batch_size, src_len, embed_dim)``.
	key_padding_mask (Tensor or None, optional): Mask to exclude keys that are pads, of shape
	`(batch, src_len)`, where padding elements are indicated by 1s. (Default: ``None``)
	attn_mask: Needs to be ``None``. The argument exists for compatibility with
	``EncoderLayer``. (Default: ``None``)
	position_bias (Tensor or None, optional): Position bias of shape
	``(batch_size * num_heads, src_len, src_len)``. When used inside WavLM model encoder, will be
	generated in the first layer and then passed from each encoder layer to the next one.
	(Default: ``None``)
	Returns:
	attn_output (Tensor): Attention output of shape ``(batch_size, src_len, embed_dim)``.
	position_bias (Tensor or None): Position bias of shape ``(batch_size * num_heads, src_len, src_len)``.
	"""
	bsz, seq_len, embed_dim = query.size()
	assert embed_dim == self.embed_dim
	assert key_padding_mask is None

	# only for the first layer
	if self.rel_attn_embed is not None and position_bias is None:
	position_bias = self.compute_bias(seq_len, seq_len)
	position_bias = position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1).view(bsz * self.total_num_heads, seq_len, seq_len)

	attn_mask_rel_pos: Optional[Tensor] = None
	if position_bias is not None:
	attn_mask_rel_pos = position_bias
	if self.gru_rel_pos: # Apply gating on relative position bias
	query_layer = query.view(bsz, seq_len, self.total_num_heads, -1)
	query_layer = query_layer.permute(0, 2, 1, 3)

	gate_a, gate_b = torch.sigmoid(
	self.gru_rel_pos_linear(query_layer).view(bsz, self.total_num_heads, seq_len, 2, 4).sum(-1, keepdim=False)
	).chunk(2, dim=-1)
	gate_a_1 = gate_a * (gate_b * self.gru_rel_pos_const - 1.0) + 2.0
	attn_mask_rel_pos = gate_a_1.view(bsz * self.total_num_heads, -1, 1) * position_bias

	attn_mask_rel_pos = attn_mask_rel_pos.view((-1, seq_len, seq_len))
	attn_mask_rel_pos = attn_mask_rel_pos.reshape(bsz, self.total_num_heads, seq_len, seq_len)[:, self.remaining_heads, :, :]

	attn_mask = attn_mask_rel_pos
	if attention_mask is not None:
	attn_mask = attn_mask + attention_mask
	if key_padding_mask is not None:
	attn_mask = attn_mask.masked_fill(
	key_padding_mask.reshape(bsz, 1, 1, seq_len),
	float("-inf")
	)
	attn_output, _ = super().forward(query, attention_mask=attn_mask)

	return attn_output, position_bias

	def prune(self):
	new_config = {
	"use_attention": True,
	"remaining_heads": self.remaining_heads,
	}
	if self.hard_concrete_for_layer is not None:
	assert not self.hard_concrete_for_layer.training
	layer_mask = self.hard_concrete_for_layer() # (1,)
	self.out_proj.weight.data *= layer_mask
	self.out_proj.bias.data *= layer_mask
	if layer_mask == 0:
	new_config["use_attention"] = False
	self.hard_concrete_for_layer = None

	if self.hard_concrete_for_heads is not None:
	assert not self.hard_concrete_for_heads.training
	head_mask = self.hard_concrete_for_heads() # (num_heads,)
	new_config["remaining_heads"] = head_mask.nonzero().squeeze(-1).tolist()
	if len(new_config["remaining_heads"]) == 0:
	new_config["use_attention"] = False
	else:
	full_mask = head_mask.repeat_interleave(self.head_dim)
	full_index = full_mask.nonzero().squeeze(-1) # 1D

	prune_linear_layer(self.k_proj, full_index, "output")
	prune_linear_layer(self.v_proj, full_index, "output")
	prune_linear_layer(self.q_proj, full_index, "output")

	self.out_proj.weight.data *= full_mask
	prune_linear_layer(self.out_proj, full_index, "input")
	self.hard_concrete_for_heads = None

	return new_config


	class FeedForward(Module):
	"""Layer that follows attention layer in encoder layer."""

	def __init__(
	self,
	io_features: int,
	intermediate_features: int,
	intermediate_dropout: float,
	output_dropout: float,
	prune_intermediate: bool = False,
	prune_layer: bool = False,
	):
	super().__init__()
	self.intermediate_dense = nn.Linear(io_features, intermediate_features)
	self.intermediate_dropout = nn.Dropout(intermediate_dropout)
	self.output_dense = nn.Linear(intermediate_features, io_features)
	self.output_dropout = nn.Dropout(output_dropout)

	if prune_intermediate:
	self.hard_concrete_for_intermediate = HardConcrete(
	n_in=intermediate_features, init_mean=0.5
	)
	else:
	self.hard_concrete_for_intermediate = None

	if prune_layer:
	self.hard_concrete_for_layer = HardConcrete(n_in=1, init_mean=0.01)
	else:
	self.hard_concrete_for_layer = None

	def forward(self, x):
	"""
	Args:
	x (Tensor): shape: `(batch, sequence_length, io_features)`
	Returns:
	x (Tensor): shape: `(batch, sequence_length, io_features)`
	"""
	x = self.intermediate_dense(x)
	x = torch.nn.functional.gelu(x)
	x = self.intermediate_dropout(x)

	if self.hard_concrete_for_intermediate is not None:
	intermediate_mask = self.hard_concrete_for_intermediate() # (intermediate_features,)
	x = x * intermediate_mask

	x = self.output_dense(x)
	x = self.output_dropout(x)

	if self.hard_concrete_for_layer is not None:
	layer_mask = self.hard_concrete_for_layer() # (1,)
	x = x * layer_mask

	return x

	def get_num_params(self):
	io_features = self.intermediate_dense.in_features
	if self.hard_concrete_for_intermediate is not None:
	intermediate_features = self.hard_concrete_for_intermediate.l0_norm()
	else:
	intermediate_features = self.intermediate_dense.out_features
	num_params = (io_features + 1) * intermediate_features + (intermediate_features + 1) * io_features

	if self.hard_concrete_for_layer is not None:
	num_params *= self.hard_concrete_for_layer.l0_norm()

	return num_params

	def prune(self):
	new_config = {
	"use_feed_forward": True,
	"ff_interm_features": self.intermediate_dense.out_features
	}
	if self.hard_concrete_for_layer is not None:
	assert not self.hard_concrete_for_layer.training
	layer_mask = self.hard_concrete_for_layer()
	self.output_dense.weight.data *= layer_mask
	self.output_dense.bias.data *= layer_mask
	if layer_mask == 0:
	new_config["use_feed_forward"] = False
	self.hard_concrete_for_layer = None

	if self.hard_concrete_for_intermediate is not None:
	assert not self.hard_concrete_for_intermediate.training
	interm_mask = self.hard_concrete_for_intermediate()
	interm_index = interm_mask.nonzero().squeeze(-1) # NOTE: must specify dim=-1
	new_config["ff_interm_features"] = len(interm_index)
	if new_config["ff_interm_features"] == 0:
	new_config["use_feed_forward"] = False
	else:
	prune_linear_layer(self.intermediate_dense, interm_index, "output")

	self.output_dense.weight.data *= interm_mask
	prune_linear_layer(self.output_dense, interm_index, "input")
	self.hard_concrete_for_intermediate = None

	return new_config


	class EncoderLayer(Module):
	"""A layer unit in encoder. Combines multihead self attention and feed forward."""

	def __init__(
	self,
	attention: Optional[Module], # can be None if the entire layer is pruned
	dropout: float,
	layer_norm_first: bool,
	feed_forward: Optional[Module], # can be None if the entire layer is pruned
	embed_dim: int,
	):
	super().__init__()
	self.attention = attention
	self.dropout = nn.Dropout(dropout)
	self.layer_norm = nn.LayerNorm(embed_dim)
	self.layer_norm_first = layer_norm_first
	self.feed_forward = feed_forward
	self.final_layer_norm = nn.LayerNorm(embed_dim)
	self.embed_dim = embed_dim

	def forward(
	self,
	x: Tensor,
	attention_mask: Optional[Tensor] = None,
	position_bias: Optional[Tensor] = None,
	key_padding_mask: Optional[Tensor] = None,
	) -> Tuple[Tensor, Optional[Tensor]]:
	"""
	Args:
	x (Tensor): Input of shape ``(batch, sequence_length, embed_dim)``.
	attention_mask (Tensor or ``None``, optional): attention mask
	of shape ``(batch, 1, sequence_length, sequence_length)``. (Default: ``None``)
	position_bias (Tensor or ``None``, optional): position bias of shape
	``(batch_size * num_heads, src_len, src_len)``.
	Only necessary for WavLM model, ``None`` otherwise. (Default: ``None``)
	key_padding_mask (Tensor or ``None``, optional): key padding mask of shape ``(batch_size, src_len)``.
	Only used for WavLM model, ignored otherwise. (Default: ``None``)
	Returns:
	(x, position_bias): Shapes are the same as in the input. Position bias is only relevant for WaLM model,
	``None`` otherwise.
	"""
	if self.attention is not None:
	residual = x

	if self.layer_norm_first:
	x = self.layer_norm(x)

	x, position_bias = self.attention(
	x, attention_mask=attention_mask, position_bias=position_bias, key_padding_mask=key_padding_mask
	)

	x = self.dropout(x)
	x = residual + x

	if self.layer_norm_first:
	if self.feed_forward is not None:
	x = x + self.feed_forward(self.final_layer_norm(x))
	else:
	# NOTE: for post norm, the layer norms should always be applied even if the layers are pruned.
	x = self.layer_norm(x)
	if self.feed_forward is not None:
	x = x + self.feed_forward(x)
	x = self.final_layer_norm(x)
	return x, position_bias

	def get_num_params(self):
	num_params = self.embed_dim * 2 * 2 # two layer norms
	if self.attention is not None:
	num_params += self.attention.get_num_params()
	if self.feed_forward is not None:
	num_params += self.feed_forward.get_num_params()
	return num_params


	class Transformer(Module):
	def __init__(
	self,
	pos_conv_embed: Module,
	dropout: float,
	layers: Module,
	layer_norm_first: bool,
	layer_drop: float,
	):
	super().__init__()
	self.pos_conv_embed = pos_conv_embed
	self.layer_norm = nn.LayerNorm(pos_conv_embed.embed_dim)
	self.layer_norm_first = layer_norm_first
	self.layer_drop = layer_drop
	self.dropout = nn.Dropout(dropout)
	self.layers = layers

	def _preprocess(self, x: Tensor):
	x = x + self.pos_conv_embed(x)

	if self.layer_norm_first:
	x = self.layer_norm(x)

	x = self.dropout(x)
	return x

	def forward(
	self,
	x: Tensor,
	attention_mask: Optional[Tensor] = None,
	position_bias: Optional[Tensor] = None,
	) -> Tensor:
	x = self._preprocess(x)
	for layer in self.layers:
	if not (self.training and torch.rand(1).item() <= self.layer_drop):
	x, position_bias = layer(x, attention_mask, position_bias=position_bias)

	if not self.layer_norm_first:
	x = self.layer_norm(x)
	return x

	def get_intermediate_outputs(
	self,
	x: Tensor,
	attention_mask: Optional[Tensor] = None,
	num_layers: Optional[int] = None,
	position_bias: Optional[Tensor] = None,
	) -> List[Tensor]:
	if num_layers is not None:
	if not 0 < num_layers <= len(self.layers):
	raise ValueError(f"`num_layers` must be between [1, {len(self.layers)}]")

	ret: List[Tensor] = []
	x = self._preprocess(x)
	for layer in self.layers:
	x, position_bias = layer(x, attention_mask, position_bias=position_bias)
	ret.append(x)
	if num_layers is not None and len(ret) >= num_layers:
	return ret
	return ret

	def get_num_params(self):
	# pos_conv_embed and layer_norm
	num_params = sum(p.numel() for p in self.pos_conv_embed.parameters()) + self.pos_conv_embed.embed_dim * 2
	for layer in self.layers:
	num_params += layer.get_num_params()
	return num_params

	def prune(self):
	new_config = defaultdict(list)
	for layer in self.layers:
	attention_config = layer.attention.prune()
	new_config["use_attention"].append(attention_config["use_attention"])
	if "remaining_heads" in attention_config:
	new_config["remaining_heads"].append(attention_config["remaining_heads"])
	else:
	new_config["num_heads"].append(attention_config["num_heads"])

	if not attention_config["use_attention"]:
	layer.attention = None

	ff_config = layer.feed_forward.prune()
	new_config["use_feed_forward"].append(ff_config["use_feed_forward"])
	new_config["ff_interm_features"].append(ff_config["ff_interm_features"])
	if not ff_config["use_feed_forward"]:
	layer.feed_forward = None

	return new_config


	class Encoder(Module):
	def __init__(
	self,
	feature_projection: Module,
	transformer: Module,
	):
	super().__init__()
	self.feature_projection = feature_projection
	self.transformer = transformer

	def _preprocess(
	self,
	features: Tensor,
	lengths: Optional[Tensor] = None,
	) -> Tuple[Tensor, Optional[Tensor]]:
	x = self.feature_projection(features)

	mask: Optional[Tensor] = None
	if lengths is not None:
	batch_size, max_len, _ = x.shape
	# create mask for padded elements and zero-out them
	mask = torch.arange(max_len, device=lengths.device).expand(batch_size, max_len) >= lengths[:, None]
	x[mask] = 0.0
	# extend the mask to attention shape and set weight
	mask = -10000.0 * mask[:, None, None, :].to(dtype=features.dtype)
	mask = mask.expand(batch_size, 1, max_len, max_len)
	return x, mask

	def forward(
	self,
	features: Tensor,
	lengths: Optional[Tensor] = None,
	) -> Tensor:
	x, mask = self._preprocess(features, lengths)
	x = self.transformer(x, attention_mask=mask)
	return x

	def extract_features(
	self,
	features: Tensor,
	lengths: Optional[Tensor] = None,
	num_layers: Optional[int] = None,
	) -> List[Tensor]:
	x, masks = self._preprocess(features, lengths)
	interm = self.transformer.get_intermediate_outputs(x, attention_mask=masks, num_layers=num_layers)
	return [x] + interm

	def get_num_params(self, in_features):
	"""Calculate the current model size."""
	feature_projection_size = self.feature_projection.get_num_params(in_features)
	transformer_size = self.transformer.get_num_params()
	return feature_projection_size + transformer_size

	def prune(self, conv_out_index):
	"""In-place pruning of submodules."""
	prune_layer_norm(self.feature_projection.layer_norm, conv_out_index)
	prune_linear_layer(self.feature_projection.projection, conv_out_index, "input")
	transformer_config = self.transformer.prune()
	return transformer_config


	################################################################################
	def _get_feature_extractor(
	norm_mode: str,
	shapes: List[Tuple[int, int, int]],
	bias: bool,
	prune_conv_channels: bool = False,
	) -> FeatureExtractor:
	"""
	Args:
	norm_mode (str):
	Either "group_norm" or "layer_norm".
	If "group_norm", then a single normalization is applied
	in the first convolution block. Otherwise, all the convolution
	blocks will have layer normalization.
	This option corresponds to "extractor_mode" from fairseq.
	Expected values are "group_norm" for Base arch, and
	"layer_norm" for Large arch.
	shapes (list of tuple of int):
	Configuration of convolution layers. List of convolution configuration,
	i.e. ``[(output_channel, kernel_size, stride), ...]``
	This option corresponds to "conv_feature_layers" from fairseq.
	Expected values are
	``[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] * 2``
	for all the architectures.
	bias (bool):
	Whether to include bias term to each convolution operation.
	This option corresponds to "conv_bias" from fairseq.
	Expected values are False for Base arch, and True for Large arch.

	See Also:
	* Original implementation
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L666-L733
	* "extractor_mode"
	- Def and base:
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L38-L45
	- Large:
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L52
	* "conv_feature_layers"
	- Def, base and large:
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L94-L100
	* "conv_bias"
	- Def and base:
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L101-L103
	- Large:
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L61
	"""
	if norm_mode not in ["group_norm", "layer_norm"]:
	raise ValueError("Invalid norm mode")
	blocks = []
	in_channels = 1
	for i, (out_channels, kernel_size, stride) in enumerate(shapes):
	normalization = None
	if norm_mode == "group_norm" and i == 0:
	normalization = nn.GroupNorm(
	num_groups=out_channels,
	num_channels=out_channels,
	affine=True,
	)
	elif norm_mode == "layer_norm":
	normalization = LayerNorm(
	normalized_shape=out_channels,
	elementwise_affine=True,
	)
	blocks.append(
	ConvLayerBlock(
	in_channels=in_channels,
	out_channels=out_channels,
	kernel_size=kernel_size,
	stride=stride,
	bias=bias,
	layer_norm=normalization,
	prune_conv_channels=prune_conv_channels,
	)
	)
	in_channels = out_channels
	return FeatureExtractor(nn.ModuleList(blocks))


	def _get_encoder(
	in_features: int,
	embed_dim: int,
	dropout_input: float,
	pos_conv_kernel: int,
	pos_conv_groups: int,
	num_layers: int,
	use_attention: List[bool],
	use_feed_forward: List[bool],
	num_heads: List[int],
	head_dim: int,
	attention_dropout: float,
	ff_interm_features: List[int],
	ff_interm_dropout: float,
	dropout: float,
	layer_norm_first: bool,
	layer_drop: float,
	prune_attention_heads: bool = False,
	prune_attention_layer: bool = False,
	prune_feed_forward_intermediate: bool = False,
	prune_feed_forward_layer: bool = False,
	) -> Encoder:
	"""
	Args:
	in_features (int): The number of input features.
	embed_dim (int):
	The dimension of embedding.
	This option corresponds to "encoder_embed_dim" from fairseq.
	Expected values are 768 for Base arch, and 1024 for Large arch.
	dropout_input (float):
	The dropout probability applied after the input feature is projected
	to ``embed_dim``.
	This option corresponds to "dropout_input" from fairseq.
	Expected values are 0.1 for both Base and Large arch.
	pos_conv_kernel (int):
	The kernel size of convolutional positional embeddings.
	This option corresponds to "conv_pos" from fairseq.
	Expected values are 128 for both Base and Large arch.
	pos_conv_groups (int):
	The number of groups of convolutional positional embeddings.
	This option corresponds to "conv_pos_groups" from fairseq.
	Expected values are 16 for both Base and Large arch.
	num_layers (int):
	The number of self attention layers in transformer block.
	This option corresponds to "encoder_layers" from fairseq.
	Expected values are 12 for Base and 24 for Large arch.
	num_heads (int):
	The number of heads in self attention layers.
	This option corresponds to "encoder_attention_heads" from fairseq.
	Expected values are 12 for Base and 16 for Large arch.
	attention_dropout (float):
	The dropout probability applied after softmax in self-attention layer.
	This option corresponds to "attention_dropout" from fairseq.
	Expected values are 0.1 for Base and 0.0 for Large arch.
	ff_interm_features (int):
	The dimension of hidden features in feed forward layer.
	This option corresponds to "encoder_ffn_embed_dim" from fairseq.
	Expected values are 3072 for Base and 4096 for Large arch.
	ff_interm_dropout (float):
	The dropout probability applied in feedforward layer.
	This option correspinds to "activation_dropout" from fairseq.
	Expected values are 0.1 for both Base and Large arch.
	dropout (float):
	The dropout probability applied at the end of feed forward layer.
	This option corresponds to "dropout" from fairseq.
	Expected values are 0.1 for Base and 0.0 for Large arch.
	layer_norm_first (bool):
	Control the order of layer norm in transformer layer and each encoder layer.
	If True, in transformer layer, layer norm is applied before features are fed
	to encoder layers. In encoder layer, two layer norms are applied before and after
	self attention.
	If False, in transformer layer, layer norm is applied after features are fed
	to encoder layers. In encoder layer, two layer norms are applied after self
	attention, before and after feed forward.
	This option corresponds to "layer_norm_first" from fairseq.
	Expected values are False for Base and True for Large arch.
	layer_drop (float):
	Probability to drop each encoder layer during training.
	This option corresponds to "layerdrop" from fairseq.
	Expected values are 0.1 for both Base and Large arch.

	See Also:
	* "encoder_embed_dim"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L49-L51
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L64
	* "dropout_input"
	- Def, base and large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L75-L78
	* "conv_pos"
	- Def, base and large
	NOTE: The description is wrong.
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L204-L207
	- Usage
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L756
	* "conv_pos_groups"
	- Def, base and large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L208-L211
	* "encoder_layers"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L46-L48
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L63
	* "encoder_attention_heads"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L55-L57
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L66
	* "attention_dropout"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L66-L68
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L60
	* "encoder_ffn_embed_dim"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L52-L54
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L65
	* "activation_dropout"
	- Def
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L69-L71
	- Base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/base_960h.yaml#L55
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/vox_960h.yaml#L55
	* "dropout"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L63-L65
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L59
	* "layer_norm_first"
	- Def and base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L91-L93
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L53
	* "layerdrop"
	- Def
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L72-L74
	- Base
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/base_960h.yaml#L54
	- Large
	https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/vox_960h.yaml#L54
	"""
	feature_projection = FeatureProjection(in_features, embed_dim, dropout_input)
	pos_conv = ConvolutionalPositionalEmbedding(embed_dim, pos_conv_kernel, pos_conv_groups)

	# Original impl
	# https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L768-L782
	encoder_layers = nn.ModuleList()
	for idx in range(num_layers):
	if use_attention[idx]:
	attention = SelfAttention(
	embed_dim=embed_dim,
	num_heads=num_heads[idx],
	head_dim=head_dim,
	dropout=attention_dropout,
	prune_heads=prune_attention_heads,
	prune_layer=prune_attention_layer,
	)
	else:
	attention = None
	if use_feed_forward[idx]:
	feed_forward = FeedForward(
	io_features=embed_dim,
	intermediate_features=ff_interm_features[idx],
	intermediate_dropout=ff_interm_dropout,
	output_dropout=dropout,
	prune_intermediate=prune_feed_forward_intermediate,
	prune_layer=prune_feed_forward_layer,
	)
	else:
	feed_forward = None
	encoder_layers.append(
	EncoderLayer(
	attention=attention,
	dropout=dropout,
	layer_norm_first=layer_norm_first,
	feed_forward=feed_forward,
	embed_dim=embed_dim,
	)
	)
	transformer = Transformer(
	pos_conv_embed=pos_conv,
	dropout=dropout,
	layers=encoder_layers,
	layer_norm_first=not layer_norm_first,
	layer_drop=layer_drop,
	)
	return Encoder(feature_projection, transformer)


	def _get_wavlm_encoder(
	in_features: int,
	embed_dim: int,
	dropout_input: float,
	pos_conv_kernel: int,
	pos_conv_groups: int,
	num_layers: int,
	use_attention: List[bool],
	use_feed_forward: List[bool],
	total_num_heads: List[int],
	remaining_heads: List[List[int]],
	num_buckets: int,
	max_distance: int,
	attention_dropout: float,
	ff_interm_features: List[int],
	ff_interm_dropout: float,
	dropout: float,
	layer_norm_first: bool,
	layer_drop: float,
	prune_attention_heads: bool = False,
	prune_attention_layer: bool = False,
	prune_feed_forward_intermediate: bool = False,
	prune_feed_forward_layer: bool = False,
	) -> Encoder:
	"""
	Construct encoder for WavLM model :cite:`chen2022wavlm`. The structure of the encoder and most of the argments are
	the same as in :py:func:`_get_encoder` so refer there for documentation. The only difference from Wav2Vec2 encoder
	is usage of `WavLMSelfAttention` instead of `SelfAttention` and two additional parameters: `num_buckets` and
	`max_distance`.
	Args:
	in_features (int): See :py:func:`_get_encoder`.
	embed_dim (int): See :py:func:`_get_encoder`.
	dropout_input (float): See :py:func:`_get_encoder`.
	pos_conv_kernel (int): See :py:func:`_get_encoder`.
	pos_conv_groups (int): See :py:func:`_get_encoder`.
	num_layers (int): See :py:func:`_get_encoder`.
	num_heads (int): See :py:func:`_get_encoder`.
	num_buckets (int): Number of buckets for relative position embedding.
	max_distance (int): Maximum distance for relative position embedding.
	attention_dropout (float): See :py:func:`_get_encoder`.
	ff_interm_features (int): See :py:func:`_get_encoder`.
	ff_interm_dropout (float): See :py:func:`_get_encoder`.
	dropout (float): See :py:func:`_get_encoder`.
	layer_norm_first (bool): See :py:func:`_get_encoder`.
	layer_drop (float): See :py:func:`_get_encoder`.

	"""
	feature_projection = FeatureProjection(in_features, embed_dim, dropout_input)
	pos_conv = ConvolutionalPositionalEmbedding(embed_dim, pos_conv_kernel, pos_conv_groups)

	# Original impl
	# https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L768-L782
	encoder_layers = nn.ModuleList()
	for i in range(num_layers):
	if use_attention[i]:
	attention = WavLMSelfAttention(
	embed_dim=embed_dim,
	total_num_heads=total_num_heads[i],
	remaining_heads=remaining_heads[i],
	dropout=attention_dropout,
	has_relative_attention_bias=(i == 0), # Position embedding is only necessary in the first layer.
	num_buckets=num_buckets,
	max_distance=max_distance,
	prune_heads=prune_attention_heads,
	prune_layer=prune_attention_layer,
	)
	else:
	attention = None
	if use_feed_forward[i]:
	feed_forward = FeedForward(
	io_features=embed_dim,
	intermediate_features=ff_interm_features[i],
	intermediate_dropout=ff_interm_dropout,
	output_dropout=dropout,
	prune_intermediate=prune_feed_forward_intermediate,
	prune_layer=prune_feed_forward_layer,
	)
	else:
	feed_forward = None
	encoder_layers.append(
	EncoderLayer(
	attention=attention,
	dropout=dropout,
	layer_norm_first=layer_norm_first,
	feed_forward=feed_forward,
	embed_dim=embed_dim,
	)
	)
	transformer = Transformer(
	pos_conv_embed=pos_conv,
	dropout=dropout,
	layers=encoder_layers,
	layer_norm_first=not layer_norm_first,
	layer_drop=layer_drop,
	)
	return Encoder(feature_projection, transformer)


	def _get_padding_mask(input: Tensor, lengths: Tensor) -> Tensor:
	"""Generate the padding mask given the padded input and the lengths Tensors.
	Args:
	input (Tensor): The padded Tensor of dimension `[batch, max_len, frequency]`.
	lengths (Tensor): The lengths Tensor of dimension `[batch,]`.

	Returns:
	(Tensor): The padding mask.
	"""
	batch_size, max_len, _ = input.shape
	mask = torch.arange(max_len, device=lengths.device).expand(batch_size, max_len) >= lengths[:, None]
	return mask


	class GradMultiply(torch.autograd.Function):
	@staticmethod
	def forward(ctx, x, scale):
	ctx.scale = scale
	res = x.new(x)
	return res

	@staticmethod
	def backward(ctx, grad):
	return grad * ctx.scale, None