Tokenizes output_features using SeqIO. Args: output_features: A mapping from field name to seqio.preprocessors.OutputFeaturesType.

Build a function that returns logits suitably normalized for top-k sampling. The returned function does many reductions over the last axis of the

Unflattens the first, flat batch*decoding dimension of a non-scalar array.

Unwraps an image produced by wrap. Where there is a 0 in the last channel for every spatial position, the rest of the three channels in that

Update value at given index and propagate changes upward.

Validates JobSet name (e.g. TPUs, VMs, jobs) to ensure compat with GKE. Raises: ValueError: If name is invalid.

Returns 'image' with an extra channel set to all 1s.

Broadcasts `value` to a canonical 4 dimensional attention bias shape. Raises: ValueError: If the shape of `value` is not 2, 3, or

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

Builds a config for a JobSet, which is a set of Jobs. https://github.com/kubernetes-sigs/jobset/blob/d49514bee57da8ac9aec2fcea06c3a13c21afeae

Builds a pathways head pod. The pod includes a head container, a proxy container and a resource manager container.

Applies bytes_to_unicode mapping in tf. TODO(markblee): Consider baking this into SentencePiece normalization rules. Reference:

Computes self._model.method(input_batch). Should be called inside pjit(). Args: method: The model method. pr

Compares expected and deployed slice-selection annotations. Both None means no topology is configured, which is considered equal. Otherwise p

Compute attention context. Args: probs: probs tensor, [batch, num_heads, target_length, source_length]. v_proj: value

Compute attention logits. Args: q_proj: query tensor, [batch, target_length, num_heads, per_head_dim]. k_proj: key te

Gathers the beam slices indexed by beam_indices into new beam array. Args: nested: NestedTensor or scalars (the latter ignored).

Returns the appropriate ConfigMap name based on telemetry setting. Args: config_type: Either "http" or "grpc" Returns:

Gets current job status.

Returns the number of partitions along each dim.

Binary search to find the largest finite int32 value for which the predicate is False. Ref: <https://github.com/google-research/t5x/blob/79998013

Returns a tree of the same structure as params where each leaf is the name of the sublearner to apply.

Looks for the user's config file in the search paths, or returns None if not found. A user config file may not exist if e.g. the user has never i

Expand the Nested Tensor `stats` and add summaries. Args: stats: A Nested Tensor, e.g., containing param norms or gradient statistics.

Creates an index map function for query/bias tensor.

Builds gcloud v1 node pool API resource. Args: credentials: Gcloud credentials used by googleapiclient.discovery. Returns: d

Returns an sequence of biases in this collection except those detected as zero. Returned biases are not guaranteed to be nonzero, but are gua

Gets the global replication count for each unique shard.

Pad logical dataset in preparation for batching. Args: dataset: The dataset to pad. global_batch_size: The size of the global phy

Imports parameters from T5X Layer Norm param dict into AXLearn RMSNorm model state. Corresponding T5X module is layers.LayerNorm Args:

Imports parameters from T5X linear param dict into AXLearn Embedding or Linear model state. Corresponding T5X modules are layers.Embed and layers

Imports parameters from T5X Transformer param dict into AXLearn TransformerLayer model state. Corresponding T5X layer is network.EncoderLayer and

Parses an axis pattern string into a hierarchical tuple structure. Converts patterns like "b t (g k) h" into ('b', 't', ('g', 'k'), 'h'). Ru

Compiles `fn` to run on the device mesh. _pjit can be used for functions with a commonly used signature and partitioning scheme. Subc

Pops element from self._current_example_list, returns None if the list is empty.

Prompts a user for a value. Args: field: The field name to display. required: If True, prompts until user provides a truthy value

Gets the repo root from CWD, or return CWD if not in a repo. Note that this repo can be arbitrary, and may not contain `ROOT_MODULE_NAME`.

Reschedules the jobset onto the appropriate tier. If we can identify that the node pool has incorrect selectors for the current scheduling

Computes segment-relative positions from segment_ids. For each segment, positions start from 0. Padding positions (segment_ids == 0) get position

Starts garbage collection (if not already started) in a separate thread.

Shared extend_step test: compares autoregressive decoding against prefill.

Tests that {init,prefill}_states + extend_step is equivalent to forward for `cfg`.

Times average execution time for fn call over num_iters after warmup.

Selects top-k experts, using configured topk_fn if available.

Handles state transitions for all jobs. The scheduler is used to determine which jobs to resume/pre-empt based on job priority. The a

Creates Vertex AI experiment name from output_dir.

Surfaces any `source_exc` as `target_exc` instead. Users can use multiple contexts to wrapping different `target_exc`, e.g: ``` with (

Build learner using the AdamW optimizer and a cosine lr schedule with linear warmup.

Returns a MaskFn that's the intersection of provided MaskFn's.

Asynchronously saves TF savables from `value_map` into `dir`. When this call returns, `value_map` can be safely mutated, but saving to `dir` will

Blend image1 and image2 using 'factor'. Factor can be above 0.0. A value of 0.0 means only image1 is used. A value of 1.0 means only image2

Build test flat-configs (without HF reference).

Builds the block map where True means the block is not fully masked. Args: mask_fn: The attention mask function. q_seq_len: Query

Common bundler flags. Keyword args will be forwarded to flag definitions.

Returns a FanAxes equivalent to this one where all fields are tuples.

Maps valid float arrays found in the inputs to the requested dtype in {float32, bfloat16}. Args: in_tree: The input values. to_dt

Checks that all elements in `x` are finite.

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

Pads data to a fixed length at the first dimension. Args: input_tensor: `Tensor` with any dimension. size: specifies the first di

Collect metrics from the device, it should be empty.

Computes a merge matrix by comparing prefixes.

Produces self-attention logit biases. Args: input_ids: A Tensor of shape [batch_size, seq_len]. segment_ids: An optio