Method from_config

penzai/models/simple_mlp.py:30–53

(
      cls,
      init_base_rng: jax.Array | None,
      feature_sizes: list[int],
      name: str = "mlp",
      activation_fn: Callable[[jax.Array], jax.Array] = jax.nn.relu,
      feature_axis: str = "features",
  )

  @classmethod
  def from_config(
      cls,
      init_base_rng: jax.Array | None,
      feature_sizes: list[int],
      name: str = "mlp",
      activation_fn: Callable[[jax.Array], jax.Array] = jax.nn.relu,
      feature_axis: str = "features",
  ) -> MLP:
    assert len(feature_sizes) >= 2
    children = []
    for i, (feats_in, feats_out) in enumerate(
        zip(feature_sizes[:-1], feature_sizes[1:])
    ):
      if i:
        children.append(pz.nn.Elementwise(activation_fn))
      children.append(
          pz.nn.Affine.from_config(
              name=f"{name}/Affine_{i}",
              init_base_rng=init_base_rng,
              input_axes={feature_axis: feats_in},
              output_axes={feature_axis: feats_out},
          )
      )
    return cls(sublayers=children)


@pz.pytree_dataclass(has_implicitly_inherited_fields=True)  # pytype: disable=wrong-keyword-args  # pylint: disable=line-too-long