def __init__(
      self,
      initializer: Callable[[jax.Array], T],
      name: str,
      value_structure: Any = _INFER_VALUE_STRUCTURE,
  ):
    """Constructs an uninitialized parameter.

    Args:
      initializer: Initializer to use.
      name: Name to use.
      value_structure: If provided, an explicit value structure that the
        initializer should return. If not provided, will be inferred by using
        `jax.eval_shape`.
    """
    self.initializer = initializer
    self.name = name
    if value_structure is _INFER_VALUE_STRUCTURE:

      def _get_structure(leaf):
        if isinstance(leaf, named_axes.NamedArrayBase):
          return shapecheck.ArraySpec(
              shape=leaf.positional_shape,
              named_shape=leaf.named_shape,
              dtype=leaf.dtype,
          )
        elif isinstance(leaf, jax.ShapeDtypeStruct):
          return shapecheck.ArraySpec(
              shape=leaf.shape,
              dtype=leaf.dtype,
          )
        else:
          raise ValueError(
              "Could not infer parameter shape from initializer! Unrecognized"
              f" leaf: {leaf}"
          )

      value_structure = jax.tree_util.tree_map(
          _get_structure,
          jax.eval_shape(self.initializer, jax.random.PRNGKey(0)),
          is_leaf=named_axes.is_namedarray,
      )
    self.value_structure = value_structure
Constructs an uninitialized parameter.

    Args:
      initializer: Initializer to use.
      name: Name to use.
      value_structure: If provided, an explicit value structure that the
        initializer should return. If not provided, will be inferred by using
        `jax.eval_shape`.
    
penzai/deprecated/v1/nn/parameters.py