Select API: bulk manipulation of configs#

The select API enables changing multiple parallel values in a single line of code using a “Beautiful Soup”-inspired fluent API.

fiddle.selectors.select(cfg, fn_or_cls=None, *, tag=None, match_subclasses=True, buildable_type=<class 'fiddle._src.config.Buildable'>, check_nonempty=None)[source]#

Selects sub-buildables or fields within a configuration DAG.

Example configuring attention classes:

select(my_config, MyDenseAttention).set(num_heads=12, head_dim=512)

Example configuring all activation dtypes:

select(my_config, tag=DType).set(value=jnp.float32)
Parameters:
  • cfg (Buildable) – Configuraiton to traverse.

  • fn_or_cls (Union[Callable[..., Any], Type[Any], None]) – Select by a given function or class that is being configured.

  • tag (Optional[TagType]) – If set, selects all attributes tagged by tag. This will return a TagSelection instead of a Selection, which has a slightly different API.

  • match_subclasses (bool) – If fn_or_cls is provided and a class, then also match subclasses of fn_or_cls.

  • buildable_type (Type[Buildable]) – Restrict the selection to a particular buildable type. Not valid for tag selections.

  • check_nonempty (Optional[bool]) – Whether to raise an error on empty selections. This will be true in the future.

Return type:

Selection

Returns:

A Selection, which is a TagSelection if tag is set, and a NodeSelection otherwise.

class fiddle.selectors.Selection[source]#

Base class for selections of nodes/objects/values in a config DAG.

__iter__()[source]#

Iterates over the selected values.

Return type:

Iterator[Any]

abstractmethod get(name)[source]#

Gets all values for a particular attribute.

Parameters:

name (str) – Name of the attribute on matching nodes.

Yields:

Values configured for the attribute with name name on matching nodes.

Return type:

Iterator[Any]

abstractmethod replace(value, deepcopy=True)[source]#

Replaces all selected nodes/objects/values with a new value.

Parameters:
  • value – Value to replace selected nodes/objects/values with.

  • deepcopy (bool) – Whether to deepcopy value every time it is set.

Return type:

None

abstractmethod set(**kwargs)[source]#

Sets attributes on nodes matching this selection.

Parameters:

**kwargs – Attributes to set on matching nodes.

Return type:

None

class fiddle.selectors.NodeSelection(cfg, fn_or_cls, match_subclasses, buildable_type)[source]#

Represents a selection of nodes.

This selection is declarative, so if subtrees / subgraphs of cfg change and later match or don’t match, a different set of nodes will be returned.

Generally this class is intended for modifying attributes of a buildable DAG in a way that doesn’t alter its structure. We do not pay particular attention to structure-altering modifications right now; please do not depend on such behavior.

class fiddle.selectors.TagSelection(cfg, tag)[source]#

Represents a selection of fields tagged by a given tag.