Inlines sequential nodes into their parents if possible. This function finds nodes that match ``child_filter`` within nodes that match ``parent_f

Inlines sequential nodes into their parents if possible. This function finds nodes that match ``child_filter`` within nodes that match ``parent_f

Inserts copies of ``value`` after each selected node. Args: value: Value to insert after each selected node. and_select: If True, sel

Inserts copies of ``value`` before each selected node. Args: value: Value to insert before each selected node. and_select: If True, s

Returns True if this is a `NamedArray` or `NamedArrayView`.

Checks if a class was wrapped in the `pytree_dataclass` decorator. Note that inheriting from a PyTree dataclass type does NOT produce another PyT

Generates a JAX PyTree key for a given field name. This can be overridden if more control over JAX key paths is needed. Args: field_na

Returns a new random key.

Orders a tree of named arrays to match the structure of another tree. This function takes two PyTrees and makes each NamedArray in ``value_tree``

Builds a matrix that maps "from" indices to their "to" indices.

Builds a matrix that maps "from" indices to their "to" indices.

Prefills the key-value caches based on a prompt. Args: model: The converted model we are running inference with. initial_cache_state: The i

Decorator for constructing a frozen PyTree dataclass. This decorator: - transforms the provided class into a frozen dataclass, - registers

Refines a selection by selecting within each selected subtree. Although selectors can already be refined by making additional calls, chained

Removes selected nodes from their parent sequence (a list or tuple). This is a convenience wrapper for ``.apply_and_inline(lambda x: ())``.

Renders a top-level selection object to its foldable representation. This function produces a rendering of either the selection itself, or its se

Renders the constructor of a Struct, with an open parenthesis.

Selects subtrees and replaces them based on relative keypaths. Convenience method that combines `at_keypaths` and `set_by_path`. Args:

Sets the value of the Variable.

Builds a function to render attributes of a struct. The resulting function will render pytree node fields in a different style. Args: fields

Flattens a PyTree exactly one level, or returns None if it's not a PyTree. Args: tree: Tree to flatten. Returns: If ``tree`` has any chi

Unflattens this tree node. See `jax.tree_util.register_pytree_with_keys_class`. This method should not be overridden by subclasses, since

Returns a color for this variable in Treescope.

Returns a new mutable copy of this variable.

Updates the value of this variable to match a frozen variable.

Returns the number of tokens in the vocabulary.

Creates a new instance of the class. Args: *args: Arguments to __init__. **kwargs: Keyword arguments to __init__. Returns:

Calls the intercepted method with the given arguments and kwargs. Args: args_and_kwargs: The positional and keyword arguments passed to the

Signature for a loss function. Args: model: The structure with parameters, usually a neural network model. state: Arbitrary state man

Standardizes a named array across the given axes.

Root-mean-square standardizes a named array across the given axes.

Runs each of the sublayers in sequence. Args: value: The input to the first sublayer. **side_inputs: The side inputs for all sublayer

Runs each of the sublayers in sequence. Args: value: The input to the first sublayer. **side_inputs: The side inputs for all sublayer

Runs each of the sublayers in sequence. Args: value: The input to the layer. **side_inputs: The side inputs for all sublayers. R

Returns the input unchanged.

Checks the structure of the value, then returns it.

Randomly drops out components of the input.

Runs each of the sublayers in order, then adds back the original input. Args: value: The input to the block. **side_inputs: Side inpu

Applies the attention mask to the input array. Args: x: The input array to mask. Usually the matrix of query-key dot products. **side

Applies the attention mask to the input array. Args: x: The input array to mask. Usually the matrix of query-key dot products. **side

Applies the attention mask to the input array. Args: x: The input array to mask. Usually the matrix of query-key dot products. **side

Runs the attention computation. Args: x: The input to the computation, which will be mapped to queries, keys, and values by the sub

Runs the caching attention computation and update the K/V cache state. When called, ``self.kv_cache_end_index`` should be filled with a scala

Abstract call method for a layer. Layers are model components that take one main input and optional side inputs, and produce a single output.

Retrieves tokens from the embedding table. Args: token_index: A named array of token indices. Axis names of this array must be disj

Retrieves tokens from the embedding table. Args: out_embeddings: The output embeddings that should be mapped to token logits. Shoul

Signature for a generalized linear operator `NamedArray` initializer. This signature attempts to make explicit all of the dimensions used by an

Runs the einsum operation.

Runs the linear operator.

Adds a learned bias to the value.

Scales its input by the scaling factor.

Calls the stacked sublayers under a `jax.lax.scan`.

Scores log-probabilities for the given inputs. Args: tokens: Array of token IDs, as an integer named array with a "seq" axis and po

Processes a new subsequence of tokens and adds them to the K/V cache. When called, the internal variables tracking the key-value cache will be

Scores log-probabilities for the given inputs. Args: inputs: Structure of input arguments, containing tokens, segment positions, an

Processes a new subsequence of tokens and adds them to the K/V cache. Args: inputs: Structure of input arguments, containing tokens, segmen

Abstract call method for a layer. Layers are submodels that take a single input and produce a single output. By convention, almost all model