A weighted tensor represents a weighted Summable value. Weight should be a tensor and is assumed to be non-negative. A weight of zero corresp

Specification of a layer parameter. This is a dataclass so that jax.tree.map does not expand it, treating it as a leaf node in pytrees. This

A dict with Tensor leaf nodes whose values should be vectorized.

Bastion job state. Attributes: status: Job status. metadata: Additional metadata.

Metadata for a bastion job.

A parameter to be optimized by an optimizer.

FanAxes describes axis indices corresponding to input, output, and batch axes. Note: axes not listed in {in,out,batch}_axis are assumed to be the

Configuration for rematerialization (remat) / checkpointing of layer computations. Rematerialization (also called checkpointing) is a memory-savi

Represents key/value projections. Fields: k_proj: [batch, source_length, num_kv_heads, per_head_dim], Projected key tensor. v_pro

A FactorizationSpec describes how to factorize a parameter's gradient. Used by AdaFactor optimizer for memory-efficient second-moment estimation

A bastion job. Attributes: spec: Job spec. state: Job state. command_proc: Optional process for the main command.

The invocation context for `Module.__call__()`. Attributes: name: The context name. Must be unique among sibling contexts. parent

An optax-style optimizer with a function to partition the inputs across devices. For new optimizers, using `UpdateTransformation` is preferred in

A lazily evaluated list of biases that are added together to get the final bias. The implementation automatically flattens nested instances of `C

Validation failure (e.g. JobSpec deserialization).

A summary that should be logged as a batch of images. The shape should either be (batch, height, width, channels) or (batch, height, width).

Quota information for job scheduling.

A causal attention mask.

Records overall training goodput.

Implicit outputs from module invocations. Usually users do not interact with `OutputCollection` directly, but via `InvocationContext.add_{sum

Captures an exponential moving average for all params.

An update to model params and state that can be transformed.

A mesh shape for hybrid (i.e., ICI and DCN) parallelism. For example, with mesh axes (data, model): - Pure fsdp on a v4-8: HybridMesh

A sliding window attention mask.

Retrieve different summaries for a git repo. Git summaries include things like the current commit, branch, origin, and whether the workspace

A table which can be pretty-printed.

A segment tree adapter for best-fit decreasing bin packing.

Defines a dataset mixture component to be used for training a LMs.

Holds per-resource-type usage and quota for a project. Attributes: usage: Current resource usage for this resource type. quota: R

A dummy vocab.

A CommandGroup groups multiple commands under a common namespace.

Output of decoding.

An attention bias represented as an implicit boolean mask.

BPE tokenization in the style of GPT2 and RoBERTa, compatible with seqio. Known differences: - Encoding is truncated at null bytes, seeming

Similar to GlobalAsyncCheckpointManager but with few improvements: 1. Tensorstore calls now run in a background event loop, hiding the cost of `t

Computable mask for splash attention that supports custom mask functions. This mask accepts any Jax/Numpy exchangeable mask function following th

Generic class to define various git summary attributes and how the values for this attributes are collected and exposed as a label and/or a file.

Represents a lifecycle event for a job. Attributes: job_name: The name of the job associated with this event. state: The state of

State returned by functions in skip_and_clip_by_global_norm().

"Attention bias that adds zero. It is better to check whether a bias has `value()` None rather than using an isinstacne check on this class,

Represents the result of a background commit.

Configures the load balancer which exposes a K8s replicated job. The jobset-controller will take care of creating the load balancer ba

Configures job resources, e.g. TPU or GPU. Attributes: instance_type: Instance type, e.g. tpu-v4-8. The format of instance type is

Information about expected logged audio summaries.

FT trainer manager gRPC client with simplified, consistent interface.

Descriptor for async copy of multiple K/V pages from HBM. Ported from https://github.com/jax-ml/jax/blob/127aa7621868cb77e552b5d1f90e4a42b09c

An attention bias defined by segment ids.

Information associated with a VM instance.

A dummy config class.

A wrapper for TrainerConfigFn that tries to serialize them before calling. This helps prevent cases where trainer config functions share states in

Audio summary. Attributes: _value: A Tensor representing audio data with shape [t,] or [t,c], t is the number of frames, and

A bias and residual where the bias has type `B` (or is None) and the residual has any type. Used to represent an original bias that has been

Data object to encapsulate download job's various outputs. Attributes: jobs: A mapping from job name to Job(spec, state). jobs_wi

Iterator for ElasticDataset

Fake libtpu server for testing.

Configures the GCS FUSE mount. https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/cloud-storage-fuse-csi-driver#sidecar-co

Configures the hostPath mount. https://kubernetes.io/docs/concepts/storage/volumes/#hostpath Attributes: host_path: Host path to mou

Represents a job that is executed by bastion.

Describes whether the job should run. Attributes: over_limits: If the job cannot be scheduled, the set of resource types on which the job

A modality-specific vocab info. The range [placeholder_start, placeholder_end) denotes the global range of placeholder tokens for this modali

State for the Lion algorithm.

A class wrapper around the splash attention kernel.

Early stopping on suffix-matches.

An attention bias represented as an explicit Tensor.

Specification of a Tensor. Used to describe model parameters and optimizer states.

An error with VM creation.

Stores information for a maybe sliced jax.Shard. Attributes: data: The actual data of the shard. index: The index of the shard.

Opaque container that is not traversed by jax.tree.map.

Alignment loop state, internal use only.

Paged KV storage emitted by `PagedKVCache` (bf16 pages). The fields `k_proj` / `v_proj` mirror `KVState`'s naming so call sites that accessed