Patches the command string. Args: command: The original user command string. Returns: The patched command st

Computes a mask over sequence positions for each given length. Args: lengths: [...]. int32 max_len: T, int dtype: outputs

Stops and disposes resources.

Generates QKV and Bias for unit test purposes.

Infer the value for -1 from len(jax.devices()) and other dims if there is -1 in mesh shape. Args: mesh_shape: The original MeshShape, whi

Recursively visits objects within this Config instance. Visit can traverse Config, lists, tuples, dataclasses, and namedtuples. By default, v

Adds a state update to the output collection of the current context.

Sets Dropout.Config.rate recursively. Args: cfg: the root config under which to look for Dropout.Config. dropout_rate: the target

Sets the environment variable DATA_DIR to the given `data_dir`. Args: data_dir: The data_dir. Raises: ValueError: If the env

Returns a causal MaskFn for sliding window attentions of a given window size. Implements the `MaskFn` protocol. Note: Setting sliding_window

Similar to jax.tree.map(), but vectorizes `fn` on VDict's.

Default settings to use in tests.

Converts a spec to a bundler. Possible options: - remote_dir: The remote directory to copy the bundle to. Must be compatible with tf_

Matches the given string against a sequence of regex-based rules. Args: x: The str to match against the rules. rules: A sequence

Sets config.bias to `bias` for all relevant descendant configs in `cfg`.

Infers tpu type (e.g. v4-8 or v6e-8-1) from instance type (e.g. tpu-v4-8, v4-8, tpu-v6e-8-1 or v6e-8-1).

Returns True if a probe type has been configured (grpc_port or http_port is set).

Split `bias` into an iterable of biases of `len(cls) + 1` instances, where the ith instances has type cls[i] or ZeroAttentionBias. Each bias

Configures how the Transformer or Conformer stack will save the linearization points. We try to save activations from the forward pass that are i

Reshape scan outputs back from [num_chunks, ...] to original shape. Trims off any padding added to make seq_len divisible by scan_chunk.

Submits the command to bastion.

A wrapper around jax.lax.dot to conduct float32 matmul

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

Build a decoder transformer config in the style of GPT. Reference: https://github.com/openai/gpt-2. Args: stack_cfg: A config of Sta

Computes groundtruth assignment for anchors based on pairwise similarity. Args: per_level_anchor_boxes: A dictionary containing a

Converts `grain.MapDataset` to `grain.IterDataset`. If the dataset is already a `grain.IterDataset`, is a no-op. See also `On grain.MapDatase

JAX implementation of `einops.repeat`. einops.repeat allows reordering elements and repeating them in arbitrary combinations. This operation

Sets normalization layer configs to `norm_cfg` for all relevant descendants in `cfg`.

Loads job spec from filepath or file.

Retrieve key/value from page tables given pages. Args: page_tables: [batch_size, pages_per_sequence], specifying page indices. pa

Replace the feature keys according to mapping in `key_map`. Like seqio's rekey, except: 1. We allow for a configurable default value

Implements the pipeline recurrence. The input `fn` specifies the single-stage implementation. Each scan iteration typically wraps `ja

Compute Top@K accuracy. Args: sim: The similarity logits between each source and target. Shape: [M, N]. dtype: float. gt_targets:

Compute output paddings w.r.t. conv_padding. The output paddings value is determined by the padding value at the anchor point in the window.

Computes contrastive logits between two tensors. Args: x: A float Tensor of shape [num_features_in_x, feature_size]. y: A float T

Flushes the written records.

Get gcloud credentials, or exits if unauthenticated. Args: impersonate_account: Service account to impersonate, if not None. impe

Loads configs for the given namespace. The strategy is as follows: 1. Attempt to read the default config file; 2. Attempt to search for a

Splits prng_key to keys iteratively and return the stacked keys. Args: prng_key: The input key. This can be a single key or a StackedKeyA

Share `module` with self's descendant modules. Args: module: The module to share. Must be a descendant module of `self`.

Find first index >= start_index where predicate is True.