Instantiates an instance of `T`. The field values specified in **kwargs take precedence over those set in the config.

Return a new bias whose dtype is `dtype`.

Convenience function for accessing box coordinates. Returns: A tensor with shape [N, 4] representing box coordinates.

Adds a child module. Args: name: The child module name. child_config: The config of the child module. **k

Returns (key, value) pairs sorted by keys.

Returns an instance of FunctionConfigBase, which invokes `fn` upon instantiation. Example: ``` cfg = config_for_function(pow).set

Convenience function for setting box coordinates. Args: boxes: a tensor of shape [N, 4] representing box corners Raises:

Returns a clone of the original config with the optional keyword overrides. Unlike `self.set`, this function does not modify the config in-pl

Computes `model_params` updates from `update`. Args: updates: The updates to potentially transform and then apply. Retur

Adds the named value to `OutputCollection.summaries`. Args: name: The name of the item to add. value: The value to ad

Flattens `tree` and returns a list of (path, value) pairs.

Builds a 6D logical mesh from the provided spec. Args: pipeline: Pipeline-paralellism. Typically means partitioning model layers across t

Returns a tree of the same structure as `nested_tensor` but with corresponding shapes instead of tensors.

Inverts a boolean mask. Commonly used to switch between paddings and mask. Args: mask: A boolean tensor. Returns: A boo

Returns a value for logging.

Reads a specific value from config file under the "GCP" namespace. Args: key: The config field. fv: The flag values, which can ov

Load golden data for a specific test from the module-level .npz file. Args: module_name: Dotted module path, e.g. "axlearn.common.causal_

Update an existing job without resubmission. This will fetch the existing job from Bastion, change the trainer command, increment the

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

Traces and lowers the function using the provided arguments.

Waits for async writes to finish.

Logs this class is unsupported with `reason`. The log message will be formatted as `Not using {self.name()} because {reason}`. This

Constructs a JobSpec with basic schema validation.

Invokes <module>.<method> in a pure functional fashion. The invocation will not depend on external inputs or have any side effects. The results o

Recursively indexes through the nested tensor. Args: x: The tensor to index. path: The sequence of keys used to recursively

Runs training. Args: prng_key: The pseudo random generator key. return_evaler_summaries: Whether to force run evalers

JAX implementation of `einops.rearrange`. This function provides a concise and readable way to perform common tensor manipulations such as re

Detokenizes int sequence to a string, up through first EOS.

Returns a tree of the same structure as `nested_tensor` but with corresponding paths instead of values. E.g., tree_paths({'a': 1, 'b'

Define flags on `fv` by recursively invoking `cfg.klass.define_flags`. Flags are defined in topological order, i.e., parent flags will be defined

Analogous to tf.io.gfile.GFile.

Read values from `fv` by recursively invoking `cfg.klass.from_flags`. The config precedence is `kwargs` followed by `fv` followed by `cfg`. In ot

Checks if a function intended for a mesh shape is compatible with the current device(s).

Writes data in a non-blocking manner.

Tags 'x' with 'name' using a custom jax.extend.core.Primitive, which is otherwise a no-op. This is useful for custom activation remat

Find the smallest available capacity >= item_weight.

Returns a fake grain input source.

Applies **kwargs to the given `cfg` if the keys exist.

Replicates and converts Tensors in `x` to local DeviceArrays. Args: x: The tensor to replicate. Returns: Replicated tensor.

Return whether to the bias has a value.

Compile the function with provided options.

Compute the cross entropy loss between logits and target_labels. Computes a stabilized-gradient version of: -jnp.sum(targets * nn.log_sof

Query operation result over range [left, right] inclusive.

Raises ValueError if any of the given `paths` are not present in `x`.

Returns an instance of ClassConfigBase, which is an object factory for `cls`. In other words, instantiating the config produces an instance of `c

Returns a debug string for the config. Args: kv_separator: The key-value separator. field_separator: The field separa

Tokenizes string to an int sequence.