<jupyter_start><jupyter_text>Fiddle `select()` and Tag APIs*Please see https://github.com/google/fiddle/blob/main/docs/colabs.md for more colabs.*The `select()` and `Tag`/`TaggedValue` APIs for Fiddle let users conciselychange many values in a larger configuration structure.The `select()` API makes it easy to set parameters across all occurrences ofspecific functions or classes within a config. For example:```python Set all Dropout classes to have rate 0.1.select(root_cfg, nn.Dropout).set(rate=0.1)```Values can also be tagged with one or more tags, making it easy to set valuesthat are shared in many places all at once. For example:```python Set all tagged dtypes, which may be on different functions/classes.select(root_cfg, tag=ActivationDType).set(value=jnp.bfloat16)```Both of these APIs facilitate configuration to be factorized into setupsinvolving a declaration of the base model (say, `base_model.py`) and severalexperiment override files (say, `my_experiment_1.py`, `my_experiment_2.py`),where the latter sets values tagged by TaggedValues, or modifies specificfunctions or classes using `select`.<jupyter_code>!pip install fiddle-config


import fiddle as fdl
from fiddle import graphviz
from fiddle import printing
from fiddle import tagging
from fiddle.experimental import selectors
import fiddle.extensions.jax

fiddle.extensions.jax.enable()  # Nicer printout for JAX types; non-essential.<jupyter_output><empty_output><jupyter_text>Running exampleLet's first consider a simple structure of Flax modules. These ones add a range(e.g. an array `[0, 1, 2, 3]`) to their input.<jupyter_code>from typing import Any

from flax import linen as nn
from jax import numpy as jnp


class AddRange(nn.Module):
  start: int
  stop: int
  dtype: Any

  def __call__(self, x):
    return x + jnp.arange(self.start, self.stop, dtype=self.dtype)


class AddTwoRanges(nn.Module):
  add_range_1: AddRange
  add_range_2: AddRange

  def __call__(self, x):
    return self.add_range_2(self.add_range_1(x))


cfg = fdl.Config(AddTwoRanges)
cfg.add_range_1 = fdl.Config(AddRange, 0, 4, jnp.float32)
cfg.add_range_2 = fdl.Config(AddRange, 0, 4, jnp.float32)
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>This model can be run as follows (see[this colab](https://colab.sandbox.google.com/github/google/flax/blob/master/docs/notebooks/linen_intro.ipynb)for an introduction to Flax APIs),<jupyter_code>model = fdl.build(cfg)
model.apply({}, jnp.array([1, 2, 1, 2]))<jupyter_output><empty_output><jupyter_text>`select()` APITo enable easier wide-spread modification of configurations, we add a simpletool to select nodes across the configuration DAG, and then set new values onthem.The main call into this method is `select()`. It currently takes a root config,and a function or class to select. It returns a `Selection` object,<jupyter_code>selectors.select(cfg, AddRange)<jupyter_output><empty_output><jupyter_text>This `Selection` object supports iteration over selected nodes,<jupyter_code>list(selectors.select(cfg, AddRange))<jupyter_output><empty_output><jupyter_text>Let's say we wanted an integer version of our model. Because the `arange` callshave a dtype (hyper)parameter, just sending integer inputs doesn't work (JAXauto-casts the integer side of the addition to a `float`),<jupyter_code># Notice that the output dtype is float32.
model.apply({}, jnp.array([1, 2, 1, 2], dtype=jnp.int32))<jupyter_output><empty_output><jupyter_text>We can use the `select()` API to set both `AddRange` dtypes to `int32`,<jupyter_code>selectors.select(cfg, AddRange).dtype = jnp.int32
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>and this will correctly have an integer output,<jupyter_code>model = fdl.build(cfg)
model.apply({}, jnp.array([1, 2, 1, 2], dtype=jnp.int32))<jupyter_output><empty_output><jupyter_text>`select()` also has a shorthand for setting multiple values,<jupyter_code>selectors.select(cfg, AddRange).set(start=1, stop=10)
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>An API to get all values for a particular field is also provided; this may beuseful for unit testing or debugging,<jupyter_code>list(selectors.select(cfg, AddRange).get('dtype'))<jupyter_output><empty_output><jupyter_text>Advanced use notesThe `Selection` object does not maintain references to the actual nodes itselects, so if the configuration is modified in the meantime, the selection willpick up any added or deleted nodes. Think of it as declarative semantics. Todemonstrate,<jupyter_code>cfg = fdl.Config(AddTwoRanges)
selection = selectors.select(cfg, AddRange)
print("Current selection:", list(selection))
cfg.add_range_1 = fdl.Config(AddRange, 0, 4, jnp.float32)
print("After adding a node:", list(selection))<jupyter_output><empty_output><jupyter_text>Suggested usage patternsFinally, for larger configuration modifications, users might find the codingpattern of binding `select` to a root config using `functools.partial` useful,since it allows modifying multiple nodes quickly.<jupyter_code>import functools


class DropoutResidualBlock(nn.Module):
  """Module that runs dropout after applying a body computation."""

  dropout: nn.Module
  body: nn.Module

  def __call__(self, x):
    residual = x
    x = self.dropout(self.body(x))
    return residual + x


# Base experiment definition, typically defined in some kind of `base_model.py`.
cfg = fdl.Config(DropoutResidualBlock)
cfg.body = fdl.Config(
    AddTwoRanges,
    fdl.Config(AddRange, 0, 4, jnp.float32),
    fdl.Config(AddRange, 0, 4, jnp.float32),
)
cfg.dropout = fdl.Config(nn.Dropout, deterministic=False)

# Experimental modifications, typically in some kind of `my_experiment.py`.
select = functools.partial(selectors.select, cfg)
select(AddRange).set(start=5, stop=9)
select(nn.Dropout).set(rate=0.2)
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>The model can be run, showing some outputs are zeroed based on RNG key,<jupyter_code>import jax.random

model = fdl.build(cfg)
inputs = jnp.array([0, 0, 0, 0], dtype=jnp.float32)
print(model.apply({}, inputs, rngs={"dropout": jax.random.PRNGKey(0)}))
print(model.apply({}, inputs, rngs={"dropout": jax.random.PRNGKey(1)}))
print(model.apply({}, inputs, rngs={"dropout": jax.random.PRNGKey(2)}))<jupyter_output><empty_output><jupyter_text>TaggedValue APIFor simple cases where we only need to set some specific attributes of a singlefunction/class, `select()` with a function/class will work great. For caseswhere the attribute we want to modify affects multiple functions/classes, e.g.`dtype`, then this could become cumbersome, because `select()` operates byclass, and some classes could name their `dtype` parameter differently.Therefore, Fiddle introduces the concept of `TaggedValue`: values that aretagged with one or more tags, and can be set all at once. Let's make our examplea little more complicated, adding a constant and a range,<jupyter_code>from typing import List


class AddConstant(nn.Module):
  value: Any
  dtype: Any

  def __call__(self, x):
    return x + jnp.array(self.value, dtype=self.dtype)


class Sequential(nn.Module):
  submodules: List[nn.Module]

  def __call__(self, x):
    for module in self.submodules:
      x = module(x)
    return x


cfg = fdl.Config(
    Sequential,
    submodules=[
        fdl.Config(AddRange, 0, 4, jnp.float32),
        fdl.Config(AddConstant, 1, jnp.float32),
    ])
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>Just to demonstrate the output of this model:<jupyter_code>model = fdl.build(cfg)
model.apply({}, jnp.array([1, 2, 1, 2], dtype=jnp.int32))<jupyter_output><empty_output><jupyter_text>Tagging values in a configurationNow, let's tag the dtypes with a tag,<jupyter_code>class ActivationDType(fdl.Tag):
  "The requested data-type for module outputs."


def base_config() -> fdl.Config[Sequential]:
  add_range = fdl.Config(AddRange, 0, 4,
                         ActivationDType.new(default=jnp.float32))
  add_const = fdl.Config(AddConstant, 1,
                         ActivationDType.new(default=jnp.float32))
  return fdl.Config(Sequential, submodules=[add_range, add_const])


graphviz.render(base_config())<jupyter_output><empty_output><jupyter_text>We can now express an override configuration, which changes both dtypes to`int32`,<jupyter_code>def experiment_config() -> fdl.Config[Sequential]:
  cfg = base_config()
  selectors.select(cfg, tag=ActivationDType).set(value=jnp.int32)
  return cfg


graphviz.render(experiment_config())<jupyter_output><empty_output><jupyter_text>and this model has `int32` output, as desired,<jupyter_code>model: Sequential = fdl.build(experiment_config())
model.apply({}, jnp.array([1, 2, 1, 2], dtype=jnp.int32))<jupyter_output><empty_output><jupyter_text>Basic API referenceIn more detail, the `TagSubclass.new()` syntax creates a `TaggedValue`, with anoptional default value.<jupyter_code>ActivationDType.new(default=jnp.float32)<jupyter_output><empty_output><jupyter_text>This is equivalent to explicitly constructing a `TaggedValue`,<jupyter_code>fdl.TaggedValue(tags={ActivationDType}, default=jnp.float32)<jupyter_output><empty_output><jupyter_text>we'll see later how to use `TaggedValue`s that have a set of tags.If you build a configuration with `TaggedValue`s that do not have a default, youwill get an error,<jupyter_code>try:
  fdl.build(ActivationDType.new())
except Exception as e:
  name = e.__class__.__name__
  %html <span style="color:red">{name}: {e}</span>
else:
  raise AssertionError("Expected an exception to be thrown!")<jupyter_output><empty_output><jupyter_text>Listing all tags in a configurationWhen configurations get very large, it can be very useful to list all availabletags. Tags sometimes serve as a high-level "API" to large configurations.<jupyter_code>for tag in tagging.list_tags(experiment_config()):
  print(tag.name, "-", tag.description)<jupyter_output><empty_output><jupyter_text>Tag subclassing and sets of tagsSince tags are types, we have a natural way of specifying a hierarchy of tags,through subclassing. We could use these to make finer-grained tags, separatingthe activation dtypes of intermediate layers and final layers.<jupyter_code>class IntermediateLayerActivationDtype(ActivationDType):
  """DType for intermediate layer neural network computations."""


class FinalLayerActivationDtype(ActivationDType):
  """DType for final layer neural network computations."""


def fine_tag_types_config() -> fdl.Config[Sequential]:
  add_range = fdl.Config(
      AddRange, 0, 4, IntermediateLayerActivationDtype.new(default=jnp.float32))
  add_const = fdl.Config(AddConstant, 1,
                         FinalLayerActivationDtype.new(default=jnp.float32))
  return fdl.Config(Sequential, submodules=[add_range, add_const])


graphviz.render(fine_tag_types_config())<jupyter_output><empty_output><jupyter_text>We can now just set the intermediate layers to have a lower precision,<jupyter_code>cfg = fine_tag_types_config()
selectors.select(
    cfg, tag=IntermediateLayerActivationDtype).set(value=jnp.bfloat16)
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>or we can achieve the same result by setting all dtypes to bfloat16, and thensetting the final ones to float32,<jupyter_code>cfg = fine_tag_types_config()
selectors.select(cfg, tag=ActivationDType).set(value=jnp.bfloat16)
selectors.select(cfg, tag=FinalLayerActivationDtype).set(value=jnp.float32)
graphviz.render(cfg)<jupyter_output><empty_output><jupyter_text>Finally, you can create tagged values that manually specify a set of tags. The`select()` API will select any `TaggedValues` that contain the tag specified.<jupyter_code>class MyTagA(fdl.Tag):
  """An example tag."""


class MyTagB(fdl.Tag):
  """Another example tag."""


def foo(a, b, ab):
  return {"a": a, "b": b, "ab": ab}


cfg = fdl.Config(
    foo,
    a=MyTagA.new(
        default=1),  # Equivalent to fdl.TaggedValue({MyTagA}, default=1).
    b=MyTagB.new(default=2),
    ab=fdl.TaggedValue({MyTagA, MyTagB}, default=3))
selectors.select(cfg, tag=MyTagB).set(value=4)
print(fdl.build(cfg))
selectors.select(cfg, tag=MyTagA).set(value=7)
print(fdl.build(cfg))<jupyter_output><empty_output><jupyter_text>Advanced use notesThere are a few things to note about the `TaggedValue` API.First is that `select(cfg, tag=)` checks tags by their class hierarchy; ifyou redefine a Tag class (e.g. by re-running the cell in colab), and havediffering Tag classes in `cfg` and `` arguments, then you will not set thevalue of any tags.Secondly, by convention, the `TaggedValue` objects are not shared, only `Tag`s.This means that if you want to set the `dtype` on a sub-network of a model, e.g.just the encoder of an encoder-decoder model, you can do so. Here we demonstratesetting it on just `add_range_2`,<jupyter_code>cfg_one_set = base_config()
selectors.select(
    cfg_one_set.submodules[1], tag=ActivationDType).set(value=jnp.bfloat16)
for submodule_cfg in cfg_one_set.submodules:
  print(submodule_cfg.dtype)<jupyter_output><empty_output><jupyter_text>Finally, we've added a little logic to the Graphviz rendering and printing, soyou can see tags in your configuration,<jupyter_code>print(printing.as_str_flattened(base_config()))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10768418671291757, 'fraction_numerical': 0.010557273207148781, 'mean_word_length': 6.164775796866558, 'pattern_counts': {'":': 6, '<': 115, '<?xml version=': 0, '>': 118, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '28', 'dir': 'jupyter-structured-clean-dedup', 'id': '519331', 'license': 'Apache-2.0', 'n_tokens_mistral': 4447, 'n_tokens_neox': 4253, 'n_words': 1216, 'path': 'fiddle/examples/colabs/select_and_tag_apis.ipynb', 'repo_name': 'saeta/fiddle', 'size': '13261'}