code: def core_remat_static(
      fn,
      variables=True,
      rngs=True,
      concrete=False,
      prevent_cse=True,
      static_argnums=(),
      policy=None,
  ):
      """Flax functional core remat version with static_argnums."""
  
      static_argnums = tuple(sorted(static_argnums))
  
      def _repack_remat_args(dyn_args, static_args):
          """Remake arg list from static and dynamic args given static_argnums."""
          args = []
          s_cnt, d_cnt = 0, 0
          for i in range(len(dyn_args) + len(static_args)):
              if i in static_argnums:
                  args.append(static_args[s_cnt])
                  s_cnt += 1
              else:
                  args.append(dyn_args[d_cnt])
                  d_cnt += 1
          return tuple(args)
  
      def inner(
          scope_fn, repack_fn, variable_groups, rng_groups, *args
      ):
          static_args = tuple(
              x for i, x in enumerate(args) if i in static_argnums
          )
          dyn_args = tuple(
              x for i, x in enumerate(args) if i not in static_argnums
          )
  
          @functools.partial(
              jax.remat,
              concrete=concrete,
              prevent_cse=prevent_cse,
              policy=policy,
          )
          @functools.wraps(fn)
          def rematted(variable_groups, rng_groups, *dyn_args):
              args = _repack_remat_args(dyn_args, static_args)
              scope = scope_fn(variable_groups, rng_groups)
              y = fn(scope, *args)
              return y, repack_fn(scope)
  
          return rematted(variable_groups, rng_groups, *dyn_args)
  
      return flax.core.lift.pack(
          inner, (variables,), (variables,), (rngs,), name="remat"
      )
  
imports: import functools
  import re
  from typing import Any, Callable, Mapping, Optional, Tuple
  import flax
  from flax import linen
  from flax import struct
  from flax.core.frozen_dict import freeze
  from flax.core.frozen_dict import unfreeze
  from flax.core.lift import In
  from flax.core.lift import Out
  from flax.linen.spmd import _axis_rules
  from flax.linen.spmd import _AxisRules
  from flax.linen.spmd import _is_logical_spec
  from flax.linen.spmd import _with_sharding_constraint
  from flax.linen.spmd import Array
  from flax.linen.spmd import ArrayPytree
  from flax.linen.spmd import get_logical_axis_rules
  from flax.linen.spmd import logical_axis_rules
  from flax.linen.spmd import logical_to_mesh
  from flax.linen.spmd import logical_to_mesh_axes
  from flax.linen.spmd import LogicalPartitionSpec
  from flax.linen.spmd import LogicalPartitionSpecPytree
  from flax.linen.spmd import LogicalRules
  from flax.linen.spmd import PartitionSpecPytree
  from flax.linen.spmd import RulesFallback
  from flax.linen.spmd import set_logical_axis_rules
  from flax.linen.spmd import with_logical_constraint
  from flax.traverse_util import flatten_dict
  from flax.traverse_util import unflatten_dict
  import jax
label_desc: Use image node_id 11 for calling a global function with example usage: core_remat_static(fn, variables, rngs, concrete, prevent_cse, static_argnums, policy) and returns: flax, tuple, rematted, y,  repack_fn