code: def froze_param_optax(
      params: optax.Params,
      tx: optax.GradientTransformation,
      frozen_patterns: Optional[List[str]] = None,
      not_frozen_patterns: Optional[List[str]] = None,
  ) -> optax.GradientTransformation:
      r"""change optax optimizer to not optimize frozen parameter.
  
    Args:
      params: model parameter in FrozenDict (Tree)
      tx: optax GradientTransform function
      frozen_patterns: a list of re patterns for frozen parameter
      not_frozen_patterns: a tupe of re patterns for trainable parameter with
        scale
  
    Returns:
      optax GradientTransform function that mask out frozen parameters.
  
    Example:
      class Encoder(nn.Module):
        @nn.compact
        def __call__(self, x, train=True):
          h = nn.Dense(8, name='dense_init')(x)
          return nn.Dense(1, name='dense_output')(h)
  
      ......(after model initialization)
      tx = optax.adamw(learning_rate=1e-3, weight_decay=1e-2)
      tx = froze_param_optax(params=params, tx = tx,\
          frozen_patterns=["params/dense_init/.*"])
      opt_state = tx.init(params)
      ......
    """
      if frozen_patterns and not_frozen_patterns is None:
          frozen_masks = bv_utils.make_mask_trees(
              params, frozen_patterns
          )
          frozen_mask = jax.tree_map(
              lambda *bools: any(bools), *frozen_masks
          )
          not_frozen_mask = jax.tree_map(operator.not_, frozen_mask)
          tx = optax.chain(
              optax.masked(tx, not_frozen_mask),
              optax.masked(optax.set_to_zero(), frozen_mask),
          )
      elif not_frozen_patterns:
          not_frozen_keys = []
          scale_vals = []
          for pattern, val in not_frozen_patterns:
              if (
                  frozen_patterns is not None
                  and pattern in frozen_patterns
              ):
                  continue
              not_frozen_keys += [pattern]
              scale_vals += [val]
          not_frozen_masks = bv_utils.make_mask_trees(
              params, not_frozen_keys
          )
          not_frozen_mask = jax.tree_map(
              lambda *bools: any(bools), *not_frozen_masks
          )
          frozen_mask = jax.tree_map(operator.not_, not_frozen_mask)
          scale_txs = [
              optax.masked(optax.scale(scale_val), mask)
              for scale_val, mask in zip(scale_vals, not_frozen_masks)
          ]
          tx = optax.chain(
              optax.masked(tx, not_frozen_mask),
              optax.masked(optax.set_to_zero(), frozen_mask),
              *scale_txs,
          )
      return tx
  
imports: import dataclasses
  import operator
  import re
  from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from big_vision import utils
  from big_vision.models import vit
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax.core import frozen_dict
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.t5 import model
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 1 for calling a global function with example usage: froze_param_optax(params, tx, frozen_patterns, not_frozen_patterns) and returns: tx