def unflaxify_apply(
    module: flax.linen.Module,
    variables: flax.typing.VariableDict,
    *dummy_args,
    rngs: flax.typing.PRNGKey | flax.typing.RNGSequences | None = None,
    method: Callable[..., Any] | str | None = None,
    mutable: flax.core.scope.CollectionFilter = False,
    **dummy_kwargs,
) -> InterceptedFlaxModuleMethod:
  """Creates an `InterceptedFlaxModuleMethod` from applying a Flax module.

  Note that this function is intended for interactive exploration and to help
  migrate Flax code to Penzai. It is not intended to be used in production
  code. Not all Flax features are supported yet; in particular, transformed
  layers are not supported and have not been tested.

  Args:
    module: The flax module to apply.
    variables: A dictionary containing variables keyed by variable collections,
      with same interpretation as for `flax.linen.Module.apply`.
    *dummy_args: Positional arguments passed to the specified apply method.
      These can be arbitrary values; their purpose is to enable tracing through
      the Flax logic.
    rngs: A dict of PRNGKeys to initialize the PRNG sequences, with same
      interpretation as for `flax.linen.Module.apply`.
    method: A function to call apply on. This is generally a function in the
      module. If provided, applies this method. If not provided, applies the
      ``__call__`` method of the module. A string can also be provided to
      specify a method by name.
    mutable: Can be bool, str, or list. Specifies which collections should be
      treated as mutable: ``bool``: all/no collections are mutable. ``str``: The
      name of a single mutable collection. ``list``: A list of names of mutable
      collections.
    **dummy_kwargs: Keyword arguments passed to the specified apply method.
      These can be arbitrary values; their purpose is to enable tracing through
      the Flax logic.

  Returns:
    An intercepted version of the Flax module call, which can be manipulated
    using Penzai tools.
  """
  if rngs is None:
    rngs = {}
  original_rngs = rngs
  intercept_states_by_path: dict[Any, _FlaxModelInterceptState] = {}
  root_intercept = None

  known_states = set()
  parameter_usages = collections.defaultdict(list)
  parameter_values = {}

  def reifying_interceptor(next_fun, args, kwargs, context):
    if context.method_name == "setup":
      # Ignore setup method.
      return next_fun(*args, **kwargs)
    nonlocal root_intercept
    # We make a depth-2 copy of the
    # variable collections because Flax represents variables as nested dicts,
    # with the first key corresponding to collection name and the second to the
    # variable or submodule name. Making a copy allows us to mutate them to
    # remove variables that we've already accounted for.
    current_vars = {
        col: dict(col_vars)
        for col, col_vars in context.module.scope.variables().items()
    }
    if not intercept_states_by_path:
      # This is the root module call.
      submodule_call_path = ()
      submodule_call_name = None
    else:
      # We are intercepting a submodule call.
      assert context.module.scope.path not in intercept_states_by_path
      parent_state = intercept_states_by_path[context.module.parent.scope.path]
      assert parent_state.module is context.module.parent
      # Figure out a name for this call.
      submodule_call_name = (
          parent_state.intercept_counter,
          f"{context.module.name}.{context.method_name}",
      )
      submodule_call_path = parent_state.submodule_call_path + (
          submodule_call_name,
      )
      parent_state.intercept_counter += 1
      # Claim all of the variables/parameters owned by this module.
      for col in current_vars.keys():
        assert col in parent_state.unclaimed_collections
        if context.module.name in parent_state.unclaimed_collections[col]:
          del parent_state.unclaimed_collections[col][context.module.name]

    # Set up the state for any submodule calls.
    current_state = _FlaxModelInterceptState(
        module=context.module,
        unclaimed_collections=current_vars,
        submodule_call_path=submodule_call_path,
        intercept_counter=0,
        submodule_calls={},
    )

    # Run the submodule logic, allowing our interceptor to recursively populate
    # the state we just built.
    intercept_states_by_path[context.module.scope.path] = current_state
    the_fn_output = next_fun(*args, **kwargs)
    del intercept_states_by_path[context.module.scope.path]

    # Assemble the re-ified layer by replacing Flax variables and parameters
    # with their Penzai equivalents. Note that any variables that are owned
    # by submodules will be moved into the submodule.
    # We detect which of the variables were actually used by this method
    # by inspecting the "reservations", which Flax uses to ensure that
    # variables or submodules aren't accidentally duplicated.
    parameter_vars, variables, immutable_variables = (
        flax.core.scope.group_collections(
            current_state.unclaimed_collections, ["params", mutable, True]
        )
    )
    name_prefix = "".join(f"{s}." for s in context.module.scope.path)
    converted_parameters = {}
    for k, v in parameter_vars["params"].items():
      if context.module.scope.name_reserved(k, col="params"):
        param_name = name_prefix + k
        converted_parameters[k] = pz.nn.SharedParameterLookup(
            pz.de.SideInputRequest(pz.nn.SharedParamTag(param_name)),
            value_structure=pz.chk.as_array_structure(v),
        )
        parameter_usages[param_name].append(submodule_call_path)
        if param_name in parameter_values:
          assert parameter_values[param_name] is v
        else:
          parameter_values[param_name] = v
    converted_variables = {}
    for col, col_vars in variables.items():
      converted_variables[col] = {}
      for k, v in col_vars.items():
        if context.module.scope.name_reserved(k, col=col):
          state_name = col + ":" + name_prefix + k
          if state_name in known_states:
            converted_variables[col][k] = pz.de.SharedLocalStateRequest(
                name=state_name, category=col
            )
          else:
            known_states.add(state_name)
            converted_variables[col][k] = pz.de.FrozenLocalStateRequest(
                state=v, name=state_name, category=col
            )

    converted_immutable_variables = {}
    for col, col_vars in immutable_variables.items():
      converted_immutable_variables[col] = {}
      for k, v in col_vars.items():
        if context.module.scope.name_reserved(k, col=col):
          converted_immutable_variables[col][k] = v

    scope_data = InterceptedFlaxScopeData(
        parameters=converted_parameters,
        variables=converted_variables,
        immutable_variables=converted_immutable_variables,
        rngs={
            rng_name: pz.de.TaggedRandomRequest(tag=rng_name)
            for rng_name in original_rngs
            if context.module.scope.rng_counters[rng_name] > 0
        },
    )
    if (
        not scope_data.parameters
        and not scope_data.variables
        and not scope_data.immutable_variables
        and not scope_data.rngs
    ):
      scope_data = None

    converted_layer = InterceptedFlaxModuleMethod(
        module=context.module.clone(),
        method_name=context.method_name,
        scope_data=scope_data,
        submodule_calls=current_state.submodule_calls,
    )

    if not intercept_states_by_path:
      root_intercept = converted_layer
    else:
      # Add this call to the parent's submodule calls.
      intercept_states_by_path[
          context.module.parent.scope.path
      ].submodule_calls[submodule_call_name] = converted_layer

    return the_fn_output

  def go():
    with flax.linen.intercept_methods(reifying_interceptor):
      _ = module.apply(
          variables,
          *dummy_args,
          rngs=rngs,
          method=method,
          mutable=mutable,
          **dummy_kwargs,
      )

  jax.eval_shape(go)
  assert root_intercept is not None

  # Handle parameter sharing: we walk the tree and identify the deepest
  # ancestor that dominates all usages of the same parameter.
  params_by_owner_path = collections.defaultdict(list)
  param_singletons = set()
  for param_name, paths in parameter_usages.items():
    if len(paths) == 1:
      param_singletons.add(param_name)
    else:
      params_by_owner_path[_common_prefix(paths)].append(param_name)

  def bind_params(
      intercepted: InterceptedFlaxModuleMethod, path: tuple[str, ...]
  ) -> pz.Layer:
    # Insert parameters that should live here.
    scope_data = intercepted.scope_data
    if scope_data is not None:
      for local_name, stub in list(scope_data.parameters.items()):
        param_name = stub.ref.tag.name  # pylint: disable=attribute-error
        if param_name in param_singletons:
          scope_data.parameters[local_name] = pz.nn.Parameter(
              name=param_name, value=parameter_values[param_name]
          )
    # Update submodule calls recursively.
    for call_name, call in list(intercepted.submodule_calls.items()):
      intercepted.submodule_calls[call_name] = bind_params(
          call, path + (call_name,)
      )
    if path in params_by_owner_path:
      # Take ownership of parameters that were used across multiple submodule
      # calls at this level.
      return pz.de.WithConstantSideInputs.handling(
          body=intercepted,
          side_inputs={
              pz.nn.SharedParamTag(name): pz.nn.Parameter(
                  name=name, value=parameter_values[name]
              )
              for name in params_by_owner_path[path]
          },
          handler_id=pz.de.infer_or_check_handler_id(
              "flax_shared_params", intercepted
          ),
      )
    else:
      return intercepted

  return bind_params(root_intercept, ())
Creates an `InterceptedFlaxModuleMethod` from applying a Flax module.

  Note that this function is intended for interactive exploration and to help
  migrate Flax code to Penzai. It is not intended to be used in production
  code. Not all Flax features are supported yet; in particular, transformed
  layers are not supported and have not been tested.

  Args:
    module: The flax module to apply.
    variables: A dictionary containing variables keyed by variable collections,
      with same interpretation as for `flax.linen.Module.apply`.
    *dummy_args: Positional arguments passed to the specified apply method.
      These can be arbitrary values; their purpose is to enable tracing through
      the Flax logic.
    rngs: A dict of PRNGKeys to initialize the PRNG sequences, with same
      interpretation as for `flax.linen.Module.apply`.
    method: A function to call apply on. This is generally a function in the
      module. If provided, applies this method. If not provided, applies the
      ``__call__`` method of the module. A string can also be provided to
      specify a method by name.
    mutable: Can be bool, str, or list. Specifies which collections should be
      treated as mutable: ``bool``: all/no collections are mutable. ``str``: The
      name of a single mutable collection. ``list``: A list of names of mutable
      collections.
    **dummy_kwargs: Keyword arguments passed to the specified apply method.
      These can be arbitrary values; their purpose is to enable tracing through
      the Flax logic.

  Returns:
    An intercepted version of the Flax module call, which can be manipulated
    using Penzai tools.
  
penzai/deprecated/v1/toolshed/unflaxify.py