Builds a config for a single container. Returns: A nested dict corresponding to a k8s Container config.

Builds a config for an HTTPRoute. Returns: A nested dict corresponding to a K8s HTTPRoute config.

Check that `tree` is masked correctly. Args: tree: The tree to check. rule: The rule from `cf

Check if current role can perform action.

Computes zero-shot classification metrics. Args: text_labels: A list of class labels. text_embeddings: A list of embeddings, each

Build gcloud compute v1 API resource. Args: credentials: gcloud credentials used by googleapiclient.discovery. Returns: disco

Decodes token IDs back to a string. Args: ids: A tf.Tensor of shape [num_tokens] or [batch, num_tokens] and dtype tf.int32.

Looks up a default config file from package, repo, or relative to CWD.

Retrieves launcher by matching `action` and flags. If there are multiple matches, the first one in the registry is returned.

Execute a gRPC call with uniform debug logging. Args: hostname: Target hostname. port: Target port. metho

Produces `input_ids` autoregressively from `target_labels`. NOTE: this performs a conversion from `grain.MapDataset` to `grain.IterDataset`.

Monitor cumulative goodput if enabled. Instantiate ml-goodput-measurement's GoodputMonitor to asynchronously calculate Goodput, Badpu

Monitor rolling window goodput if enabled.

Invokes 'fn' for each sub-layer. Note, the number of sub-layers used for the computation might be smaller than `cfg.num_layers` depen

Asynchronously serializes a list of tensors with _async_serialize.

Cleans up the process/fds and upload logs to gs.

Adafactor optimizer. References: https://arxiv.org/abs/1804.04235 https://github.com/deepmind/optax/blob/c4a4790b85ad69cda00a425cc3dcf9c9

Add parameter scaled by `weight_decay * (learning_rate ** learning_rate_exponent)`. Unlike optax.add_decayed_weights, this supports scaling weigh

Adds a flag to the command group. Args: name: Name of the flag. Should include "--". undefok: If True, the flag will

A decorator that creates a wrapper around `fn` that annotates the call with annotations given by the supplied keyword arguments. These annota

Performs the model analysis on the AOT compiled JAX program. Refer to https://docs.jax.dev/en/latest/jax.stages.html#jax.stages.Compiled Not

Pads or truncates ids so as to divide max length, then group into temporary batch.

Returns a function that generates a tf.data.Dataset object. Note: batch(is_training=True) requires sufficient number of examples per host. Wh

Computes Area Under the Receiver Operating Characteristic Curve (ROC AUC) for binary classification model. Reference: https://github.com/

Assert that the trainer config debug string matches the golden file.

Computes confusion matrix. References: https://github.com/scikit-learn/scikit-learn/blob/baf0ea25d6dd034403370fea552b21a6776bef18/sklearn/met

Returns dilated effective window size. Args: window: convolution window. dilation: convolution dilation. Returns: Th

Returns the explicit padding for "SAME", "VALID", and "CAUSAL" modes. Each mode follows the formulas below: * SAME: (pad_total//2, pad_total

Analogous to tf.io.gfile.copy.

Sets target[path] = source[path]. Args: source: The source tree. target: The target tree. path: The sequence of keys used

Count the number of parameters in a model.

Creates node pools for the job.

Creates MethodRunner for the specified method and arguments. Args: method: the method name of self.model to invoke. The method sh

Computes relative position buckets as seen in DeBERTa. Reference: https://github.com/microsoft/DeBERTa/blob/771f5822798da4bef5147edfe2a4d0e82

Detokenizes int32 iterable to a string, up through first EOS.

Constrains all leaf values in the input batch, then (optionally) dispatches examples to a subset along the batch axis. The dispatchings are a

Makes evaler configs from the given input sources. Args: input_source_configs: A dictionary with the eval dataset name(s) as key(s) and

The FLOPs loss in 'Minimizing FLOPs to learn efficient sparse representations' ICLR2020. a_j = mean(|a_{ij}|, axis=0) L = sum(a_j^2) Re

Constructs a single combined attention bias of the same type as this class from a sequence of such biases. If the sequence is empty,

Get the feature map function and its backward function.

Infer MoE outer batch size from mesh shape. Args: mesh_axis_names: The name of each mesh axis. outer_batch_axis_names: The names

Initializes cache for autoregressive cached decoding. Args: batch_size: KV cache batch size. max_len: KV cache max le

Returns partition spec for the input batch. We partition the inputs along all axes. For example, if the mesh has shape (64, 4) and axis names

Return True if git diff or git porcelain are not empty. Raises: NotImplementedError if called on invalid git summary.

Ensures job name is not path-like and only contains safe characters. This check should avoid making assumptions about the underlying compute envi

Loss based on https://arxiv.org/pdf/1801.09414.pdf. Assumes that `logits` are cosine distances between examples and class embeddings. Args:

Codebook look up with ids. Args: ids: integer tensor of shape [..., num_codebooks] with values in range [0, codeb

Generates attention logit biases given the segment ids. ... such that positions belonging to different segments cannot attend to each other.

Analogous to tf.io.gfile.makedirs.

Mocks out VM get, create, and delete.

Converts time in milliseconds to number of samples under the given sample rate. Args: ms: Time in milliseconds. sample_rate: Samp

Parses action from argv, or exits with usage info. The action is inferred from the first positional arg in argv[1:] (where argv[0] is interpreted

Computes pairwise loss among ranked docs (ranks > 0). For every pair of docs (a, b) with 0 < rank_a < rank_b, i.e., (a, b) represent a pair o

Starts a save to the given directories. The writes may happen in the background and not finish until wait_until_finished(). Storage i

Scale updates by rms of the gradient for each param vector or matrix. A `block` is here a weight vector (e.g. in a Linear layer) or a weight matr

Scale the 2D image using the autocontrast rule.

Returns whether each parameter should be updated with the optimizers. Args: model_params: A nested structure with OptParams as le

Similar to `input_tf_data.unbatch`. Unlike `batch`, which naively groups top-level elements, unbatch applies to JAX leaves only. For example,

Writes configs to a file, with a prepended comment warning users not to modify it. Args: config_file: Output file path. configs:

Initializes a LoadBalancer instance. Args: jobset_name: Name of the jobset. replicated_job_name: Name of the

Builds the failurePolicy dict for a JobSet spec. Args: replicated_jobs: List of replicated job dicts (each must have a "name" key

Builds the base container with common elements across all GPU jobs Returns: A nested dict corresponding to a k8s Container config

Builds a config for the uploader container which sync logs to the output dir. The sidecar container runs an loop to periodically sync outputs

Builds a config for volumes.

Returns whether the attention kernel supports the given block size. Args: input_batch: A dict contains input entries, see __call_

Computes mean and variance over sequence data. Args: x: inputs tensor of shape [batch_size, seq_len, ...]. segment_ids: An int Te

Equivalent to `x.at[idx].get(unique_indices=unique_indices)`, but with a gather-based backward pass. The reason to use this function is that

Runs forward pass on the linear layers and dispatching and combining.

Returns the size of the dataset, by counting examples if necessary.