code: def swap_collection(fn: Callable[..., Any], col_a: str, col_b: str):
      """Swap two collections."""
  
      def swap(target):
          a = target[col_a] if col_a in target else {}
          b = target[col_b] if col_b in target else {}
          target[col_b], target[col_a] = a, b
          return target
  
      return map_variables(fn, (col_a, col_b), swap, swap, mutable=True)
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 8 for calling a global function with example usage: swap_collection(fn, col_a, col_b) and returns: map_variables, target