The default way this summary should be accumulated. Args: other: The summary from later in training/eval to accumulate into this

Converts `x` to numpy ndarray recursively. Args: x: a jnp array, numpy array, TF/PyTorch Tensor, or a nested structure of arrays or Tenso

Yields per-feed physical input batches (using `input_dispatcher` if configured). The caller should use `host_to_global_array` to construct a

Split this bias into a bias of type `cls` and a residual. If the two returned biases are added together, the result is equivalent to

Obtains the checkpoint dir for the given step.

Decodes encoded boxes with respect to anchors. Args: encoded_boxes: A [..., 4] float tensor with encoded boxes. ancho

Tokenizes string Scalar to an int32 Tensor, without adding EOS.

Expands each VDict in `tree` to a list. Args: tree: A nested tree of Tensors. All leaf nodes under a VDict must be tensors with the same

Attempts to read resource limits, both total and per-project. Also reads user quota project membership. Args: path: Absolute path to

Infer TPU version from the TPU type. Args: tpu_type: A string of the format {version}-{cores}. Returns: Inferred TPU version

List files on GS with <url> prefix. Args: url: gs:// prefixed url. Returns: List of full paths to objects with <url> prefix.

Returns runner config(s) optionally filtered by name.

Computes the output shape given input shape. Args: input_shape: Values for the input dimensions [batch_size, seq_len]. Each value

Returns a shallow copy of the input tree with subtrees pruned based on `should_prune`. This is a shallow copy because leaf nodes (non-dict values

Extracts per-param settings for the given trainer config. Given a trainer config specified by `module` and `config_name`, initializes the trainer

Registers a custom config field validator. Args: match_fn: A function that returns True if the value should be validated by `validate_fn`

Repeats key or value heads dim to match the query. TODO(dhwang2): optimize computation like GroupedQueryAttention.

Right shift the tensor and pad zeros on the left side, e.g., [2, 5, 3] with shift_size=2 -> [0, 0, 2]. Args: x: input tensor of

See `BaseCheckpointer.save` for details. In addition to behavior in `BaseCheckpointer`, saving only happens if the configured checkpo

Sets LayerNorm.Config.eps recursively. Args: cfg: The root config under which to look for LayerNorm.Config. eps: The target value

Returns the number of slices.

Updates prefix merger state given the next candidate token.

Update when a bin's capacity changes.

Waits for async CloudBuild to finish by polling for status. Is a no-op if `cfg.is_async` is False. Args: name: Bundle na

Computes the RMS norm for each leaf tensor of `x` and optionally adds summaries. Summaries will be added if `summary_suffix` is not None *and* th

Applies sharding constraints in-place for relevant named tensors in the input batch.

Executes a command on an existing VM. Args: cmd: Command to run. detached_session: If not None, run commands behind `

Constructs a filter string to query build requests based on image name and tags. To filter builds by multiple tags, use "AND", "OR", or "NOT" to

Computes start of word indexes. Note: The function expects morphun tokenized input. Args: inputs: int Tensor [seq_len] of piece ids.

Defines all launch flags, and amends `app.usage` with additional launch help info.

Wraps absl_main by returning (returncode, stdout, stderr).

Shared backward test: compares gradients of FlashAttention against reference.

Shared forward test: compares FlashAttention output against reference.

Helper to create a TPUReplicatedJob config for composite job tests.

Return a dense `(k_proj, v_proj)` pair from a dense `KVState`. Subclasses that emit paged or otherwise non-dense storage override this

Asymmetric contrastive loss from logits. When soft_labels is None, minimizing the loss will maximize the gap between logits[i, i] and logits[

Builds sharding specifications for the flash attention operation. This method is called to build the sharding map specifications that will be

Produces a code bundle for the local directory. Args: name: Bundle name. Returns: The bundle identifier.

Converts lists to delimited strings.

Produce a lowered and compiled training step. Args: trainer_state: The global trainer state (or state specs). If

Copies OptStateSpec and optionally assigns with a different memory kind. Args: specs: Nested[OptStateSpec] to copy from. pattern:

Create a timestamp protobuf with current time. Returns: Timestamp protobuf set to current time

Builds a default mapping from tensor dims to partition specs for the FlashAttention outputs. Maps attention heads over the default tensor-paralle

Default project from active `gcp_settings`. Project is used along with env_id to identify `gcp_settings`. Returns: the project in active `gc

Applies Mixup and/or Cutmix to batch of images and transforms labels. Args: images (tf.Tensor): Of shape [batch_size, height, wid

Encodes boxes with respect to anchors. Args: boxes: A [..., 4] float tensor with boxes to encode. Boxes must be of the form

Wraps _execute with retries. All args and kwargs are forwarded. Retries are triggered automatically when _execute throws an exception.

Analogous to tf.io.gfile.exists.

Wrap original function to pass in key-word args.

Gets the status of the latest CloudBuild by filtering on the build tags and image name. In order: 1. Queries the last known region where a bu

Builds a TrainerConfigFn according to the model and input specs. Args: model_cfg: The model config. learner_cfg: The learner conf

It is not thread-safe. Please see the docstring of the class for more details. Users should not call `init_nvml` multiple times.

Return True for a valid git summary, i.e. collected from valid git repository checkout, and False otherwise.

Lists running jobs and optionally prints them in tabular format. Args: output_file: Output file. If None, prints to stdout, other

Mask out tree leaves that are not transformed by the optimizer. Args: tree: A nested structure with ParameterSpec, OptParams or Tensor as

Fused streaming log-normalizer and top-k via Pallas (TPU). This kernel supports only forward. Backprop will be added latter This function is

Encode a token to id.

Returns True iff job is alive.

Produce decoder hidden states. Args: input_batch: a dict with the following entries: input_ids: an int Tensor of

Generate the (key, value) pairs for the immediate children of a pytree `node`. Reference: jax._src.tree_util.generate_key_paths() Example:

Given a layer, returns all nested parameter initialization specs. Args: layer: The layer to read from recursively. delegates: An

Restores trainer state from checkpoint. If successful, sets self._step and self._trainer_state to the restored step and state, respec

Configure the CLIP image stream encoder. Args: num_layers: An integer indicating the number of transformer blocks. model_dim: An