code: def get_optimizer(
      config: ml_collections.ConfigDict, params: PyTree
  ) -> optax.GradientTransformation:
      """Makes a Optax GradientTransformation for DeformableDETR."""
  
      # If we freeze the batch statistics for the backbone, the affine
      # transformation of a bn layer can be absorbed by its previous linear
      # transformation and therefore there is no need to train on the affine weights
      # of bn layers.
      def bn_and_freeze_batch_stats(path):
          if not config.freeze_backbone_batch_stats:
              return False
          names = ["/bn1/", "/bn2/", "/bn3/", "/init_bn/", "/proj_bn/"]
          for s in names:
              if s in path:
                  return True
          return False
  
      def early_and_load_pretrain(path):
          if not config.load_pretrained_backbone:
              return False
          names = [f"/ResidualBlock_{i}/" for i in range(3)] + [
              "/init_bn/",
              "/stem_conv/",
          ]
          for s in names:
              if s in path:
                  return True
          return False
  
      backbone_traversal = flax.traverse_util.ModelParamTraversal(
          lambda path, _: "backbone" in path
      )
      ref_embed_traversal = flax.traverse_util.ModelParamTraversal(
          lambda path, _: "ref_embed" in path
          or "sampling_offsets" in path
      )
      bn_traversal = flax.traverse_util.ModelParamTraversal(
          lambda path, _: bn_and_freeze_batch_stats(path)
      )
      early_layer_traversal = flax.traverse_util.ModelParamTraversal(
          lambda path, _: early_and_load_pretrain(path)
      )
  
      all_false = jax.tree_util.tree_map(lambda _: False, params)
  
      def get_mask(traversal):
          return traversal.update(lambda _: True, all_false)
  
      backbone_mask = get_mask(backbone_traversal)
      ref_embed_mask = get_mask(ref_embed_traversal)
      bn_mask = get_mask(bn_traversal)
      early_layer_mask = get_mask(early_layer_traversal)
  
      oc = config.optimizer_config
  
      tx = optax.chain(
          optax.clip_by_global_norm(oc.max_grad_norm),
          optax.adamw(
              learning_rate=optax.piecewise_constant_schedule(
                  oc.base_learning_rate,
                  {
                      oc.learning_rate_decay_event: oc.learning_rate_decay_rate
                  },
              ),
              b1=oc.beta1,
              b2=oc.beta2,
              weight_decay=oc.weight_decay,
          ),
          optax.masked(
              optax.scale(oc.learning_rate_reduction), backbone_mask
          ),
          optax.masked(
              optax.scale(oc.learning_rate_reduction), ref_embed_mask
          ),
          optax.masked(optax.scale(0), bn_mask),
          optax.masked(optax.scale(0), early_layer_mask),
      )
  
      return tx
  
imports: from collections.abc import Mapping
  from concurrent import futures
  import functools
  import time
  from typing import Any, Callable, Dict, Sequence, Tuple, Type, Union
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  import flax
  from flax import jax_utils
  from flax.training.checkpoints import restore_checkpoint
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.deformable_detr import coco_eval
  from scenic.projects.baselines.deformable_detr import evaluate
  from scenic.projects.baselines.deformable_detr.model import DeformableDETRModel
  from scenic.projects.baselines.detr import train_utils
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import PyTree
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 1 for calling a global function with example usage: get_optimizer(config, params) and returns: tx, False, False, traversal, False, False, True, True