A smarter version of with_sharding_constraint that only applies the constraint if the current mesh contains the axes in the partition specs.

Extracts the k from the pipeline_str.

Parition rules for GPTJ. Note that these rules are orderd, so that the beginning rules match first. It is important to use Pa

Return the global L2 norm of a pytree.

For padding the data equally. Applies padding only if tokenization_args.always_pad is true. Supports two modes of padding: * defaul

Returns a pytree of PartitionSpec according to rules. Supports handling Flax TrainState and Optax optimizer state.

An extended version of jax.tree_util.tree_map, where the mapped function f takes both the name (path) and the tree leaf as input.

Apply weight decay with schedule.

This function takes projected key, value states from a single input token and concatenates the states to cached states from previous

Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.

Make causal mask used for bi-directional self-attention.

Validate the `rope_scaling` configuration.

Basic implementation of cross-batch. If dataset_packing > 0, it assumes that the documents occupy batch entries as follows batch[0] =

For filtering the dataset according to the rules described in data_args. data_args should be separated using separate_data_args.

Generates a text file and inserts an execute line at a random position.

Return axis names from partition specs.

For getting the list of installed packages

Return a weight decay mask function that computes the pytree masks according to the given exclusion rules.

r""" Args: batch_size (`int`): batch_size used for fast auto-regressive decoding. Defines the batch size of the in