A deprecated alias for `host_to_global_array`. Please use `host_to_global_array` instead.

Return the pre-built docker image id associated with the bundler.

l2_normalize Normalizes along the dimension `axis` using an L2 norm. Args: x: Input tensor. axis: Dimension along which to normal

Returns the dimension of each attention head.

Sorts the table. Heading remains unchanged.

Validates a config field value. Validation is handled by validators registered via `register_validator`. `match_fn`s will be invoked in order

Adds tensor stats about `value`. Suppose `self.tensor_stats` adds some summaries about `value`, e.g., `self._add_tensor_stats("me

Builds a config for a single Pod, which is a set of containers. https://kubernetes.io/docs/concepts/workloads/pods Returns:

Copies contents of local directory to `temp_dir`, excluding `exclude` paths, and returns the directory. Caller is expected to use as

Add a new bin with the given capacity.

Applies `attention_logit_biases` on `logits`. Args: logits: A float Tensor. attention_logit_biases: A float Tensor. If None, assu

Like np.testing.assert_allclose, but allows outlier percentiles to be specified. `tolerance_map` is mapping of percentile values (between 0 a

See configured `decoding` implementation for details.

Builds configs for BERT Embedding layer. Defaults are from the BERT-BASE model. Args: max_position_embeddings: Number of positional

Builds a k8s liveness probe dict, omitting any None fields. Returns: A dict corresponding to a k8s liveness probe configuration.

Builds a k8s readiness probe dict, omitting any None fields. Returns: A dict corresponding to a k8s readiness probe configuration

Returns a PartitionSpec for the given partition type.

Returns the input dataset, which should produce per-feed logical batches. Each batch is a pytree of arrays which reside on host memory (i.e.,

Returns weights of shape [num_embeddings, dim].

Return the shape of the bias tensor. Note: this doesn't materialize the value. jax.eval_shape calls value(), but it only does so usin

Transforms `method_fn` (with this context) into a pure functional Callable. The returned Callable will have the same behavior as `method_fn`,

Get value at a specific index.

Computes precision, recall, and F-beta score for binary classification. References: https://github.com/scikit-learn/scikit-learn/blob/baf0ea2

Schedules jobs according to quotas. The scheduling behavior depends on the configured `cfg.scheduler`. Args: job_metadat

Writes job spec to filepath or file.

Sets x[path...] = value, where path can be a multi-part index. If any part of the path does not exist in `x`, new sub dicts will be created, e.g.

Starts status server.

Starts the async copies.

Returns a cached SentencePieceVocabulary instance to avoid repeated loading.

Checks whether a given value is appropriate for the specified summary kind. This is used to determine whether a particular summary value (or wrap

Shared gradient test: compares TPU flash attention gradients against reference.

Adds a command that invokes an existing absl script. Args: name: Command name. module: A string of the format `module

Builds configs for BERT transformer stack. Defaults are from the BERT-BASE model. Args: num_layers: Number of transformer encoder la

Tokenizes string to an int sequence, without adding EOS.

Constructs a bundler config from the given spec. Bundlers must be registered via `register_bundler`. Args: bundler_type: Type of bun

Returns workload annotations from the underlying builder.

Scales updates using a custom schedule for the step size. Unlike optax.scale_by_schedule, this implementation uses 1-based steps, i.e., the first

Selects a mesh rule (if one matches `mesh_selector` to override mesh config. If any of `trainer_config.mesh_rules` matches `mesh_selector`, modif

Create standardized error response. Args: context: gRPC context code: gRPC status code message: Error mes

See `_mock_popen_fn`.

Recursively converts defaultdicts to dicts.

Makes the local bastion state consistent with the remote state. This function serves as a synchronization point for user-initiated state chan

Handles all state transitions for a single job. Assumptions: 1. A jobspec file exists in the remote job dir at the start of each call

Split this bias into a bias of type `cls` and a residual. Compared to the default implementation, this determines which instance of `cls` to

Compute the binary cross entropy loss between logits and targets. Computes a stabilized-gradient version of: -jnp.sum(targets * jnp.log(l

Converts delimited strings to lists.

Constructs a device mesh. If `mesh_shape` is specified as a `HybridMeshShape`, we use the `ici_mesh_shape` and `dcn_mesh_shape` directly to c

Builds a default FlashAttention mapping from tensor dims to partition specs for the MHA impl. Maps attention heads over the default tensor-parall

Implements activation-kernel einsum with quantization (e.g. fakequant, fp8-fp8, etc.) Args: subscripts: Specifies the subscripts

Enable sliding window attention. Args: cfg: MultiheadAttention Config. sliding_window_size: Sliding window size. Returns:

Runs eval for the given step. Args: step: Current step. prng_key: PRNG key. model_params: Model parameter

Generates fake audio data with a fixed seed.

Generate a unique job name.

Attempts to infer the CLI name.

Makes config name string as a function of architecture and model-size. Useful to keep config names synced with fine-tuning configs. Args:

Parses sequence of k:v into a dict. Args: kv_flags: A sequence of strings in the format "k:v". If a key appears twice, the last

Recursively find the target key in the config and set its value. Args: path: A sequence of keys for indexing to set the target va

Returns a flattened dict with path -> value mappings. Args: omit_default_values: Omit a field from the output dict if its value r

Cache the fields dict for type. Args: type_obj: Type to be cached. Returns: A dictionary of fields for the type.

Paths to search for config file, ordered by precedence.

Create standardized success response. Args: response_type: Response message type **kwargs: Additional response fields

Returns a TrainerConfigFn for training ResNet on ImageNet. Args: backbone: The ResNet backbone config. train_batch_size: Global t

Gets the CloudBuild status for the latest build in a given region (no caching). Args: project_id: The GCP project ID. image_name:

Parses health check spec and returns timeout if the `check_type` is in the check_spec. Also checks if environment variables satisfy health check