x: (batch, seqlen, hidden_dim) (where hidden_dim = num heads * head dim) key_padding_mask: bool tensor of shape (batch, seqlen) Retur

Forward will take an attention mask but ignore it!

Forward will take an attention mask but ignore it!

Sequence-scaled input.

Sequence-scaled input cumulative sum.

Sequence-scaled input cumulative sum.

r"""Inputs of forward function Args: x: the sequence fed to the positional encoder model (required). Shape: x:

r"""Inputs of forward function Args: x: the sequence fed to the positional encoder model (required). Shape: x:

This is a batch-first implementation

A variety of pooling options. Some ignore the cls token. Input needs to be B S H.

This is the same eps clipping as in the original ScaleNorm implementation.

This is the same eps clipping as in the original ScaleNorm implementation.

Return default hydra config for a given attack.

Time-based scheduler with optional warmup and cooldown (so technically a trapezoidal shape)

Follows the huggingface transformers scheduler with the same name, but gets an additional arg hour_budget

Follows the huggingface transformers scheduler with the same name, but gets an additional arg hour_budget

Constant, then linear decay. plot min(1, 1/0.5 - x / (1000 * 0.5)) from 0 to 1000 in the plot range 0 to 1

An inverse cosine schedule, with limited budget.

Time-based scheduler as described in Iszak et al. plus inv_sqrt. Takes in num_warmup_steps and num_training_steps as normal, but actually squeezes

Follows the huggingface transformers scheduler with the same name, but gets an additional arg hour_budget

Simple multi-cycle scheduler. Not including paper/fastai three-phase things or asymmetry.

Simple single-cycle scheduler. Not including paper/fastai three-phase things or asymmetry.

Follows the huggingface transformers scheduler with the same name, but gets an additional arg hour_budget

to the moon.

Linear increase from a percentage of the base learning rate, then linear decay. plot min(0.5 + x * (1 - 0.5)/(1-0.25) / 1000, 1/0.25 - x / (1000

Return default hydra config.

Decay the LR based on the inverse square root of the update number. We also support a warmup phase where we linearly increase the learning rate

Return default hydra config for a given attack.

Simple single-cycle scheduler. Not including paper/fastai three-phase things or asymmetry.

to the MOON.

Load (and optionally stage) a pre-processed corpus. Create one if it doesn't exist.

This function controls the central routine.

This is boiler-plate code for a launcher.

This function controls the central routine.

This function controls the central training loop.

Load all datasets in eval.tasks for finetuning and testing.

This checkpoint can be used for downstream tasks. The default behavior is to save this checkpoint to a checkpoints folder under base_directory/nam

Save two summary tables. A detailed table of iterations/loss+acc and a summary of the end results.

Path, identifier and additional client state. This checkpoint can be used to resume training. The default behavior is to save this checkpoint rela

For cases where shapes of Q and K do not match.

(Maybe) performs a single optimization step.

Performs a single optimization step. Arguments: gradsH: The gradient used to compute Hessian vector product. closure (