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

  # Store the values for each variable.
  # This is a bit fragile. We want to call the Flax model under eval_shape so
  # that we don't need to run any compute. But it's not necessarily possible to
  # identify what the variables are until we actually identify the linen module
  # that owns each variable, since dictionaries can be nested arbitrarily.
  # This means that, when we identify the variable, we are running under
  # eval_shape.
  # To avoid creating a Penzai variable inside the eval_shape, we instead save
  # the value that that variable should have into this dictionary. If that
  # value depended on computations inside the eval_shape, this would leak a
  # tracer. But the initial value of each variable never depends on any
  # computations inside the eval_shape, so it's OK to save it here.
  # Later, outside the eval_shape, we can create the top-level pz.Variable
  # instances and substitute them in.
  pz_parameter_values = collections.defaultdict(dict)
  pz_state_var_values = collections.defaultdict(dict)

  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 = (  # pylint: disable=unbalanced-tuple-unpacking
        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 = {}
    if parameter_vars:
      for k, v in parameter_vars["params"].items():
        if context.module.scope.name_reserved(k, col="params"):
          param_label = name_prefix + k
          if param_label not in pz_parameter_values:
            pz_parameter_values[param_label] = v
          converted_parameters[k] = pz.ParameterSlot(param_label)

    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):
          var_label = FlaxVarLabel(collection=col, name=name_prefix + k)
          if var_label not in pz_state_var_values:
            pz_state_var_values[var_label] = v
          converted_variables[col][k] = pz.StateVariableSlot(var_label)

    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(  # pytype: disable=wrong-arg-types
        parameters=converted_parameters,
        variables=converted_variables,
        immutable_variables=converted_immutable_variables,
        rng_names=frozenset({
            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.rng_names
    ):
      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
  all_vars = [
      pz.Parameter(value=v, label=l) for l, v in pz_parameter_values.items()
  ] + [
      pz.StateVariable(value=v, label=l) for l, v in pz_state_var_values.items()
  ]
  return pz.bind_variables(root_intercept, all_vars)
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/toolshed/unflaxify.py