code: def test_optimizer_with_multi_lrs(self):
      config = ml_collections.ConfigDict(
          {
              "optimizer": "sgd",
              "optimizer_configs": dict(momentum=0.9),
              "lr_configs": dict(base_learning_rate=0.1),
              "layer_prefix_to_base_lrs": {
                  "video_text_fusion": 1.0,
              },
          }
      )
      expected_params = {
          "video_encoder": {
              "pos_embedding": np.zeros((8)) - 0.1,
              "block0": {
                  "kernel": np.ones((10)) - 0.1,
                  "bias": np.zeros((10)) - 0.1,
              },
          },
          "text_encoder": {
              "pos_embedding": np.zeros((8)) - 0.1,
              "block0": {
                  "kernel": np.ones((10)) - 0.1,
                  "bias": np.zeros((10)) - 0.1,
              },
          },
          "video_text_fusion": {
              "block0": {
                  "kernel": np.ones((10)) - 1.0,
                  "bias": np.zeros((10)) - 1.0,
              },
          },
      }
      tx = optimizer_utils.optimizer_with_multi_lrs(config, self.params)
      state = tx.init(self.params)
      gradients = jax.tree_util.tree_map(np.ones_like, self.params)
      updates, _ = tx.update(gradients, state, self.params)
      new_params = optax.apply_updates(self.params, updates)
      chex.assert_trees_all_close(
          flax.core.unfreeze(new_params), expected_params
      )
  
file_path: scenic/scenic/projects/unloc/optimizer_utils_test.py
imports: from absl.testing import absltest
  import chex
  import flax
  import jax
  import ml_collections
  import numpy
  import optax
  from scenic.projects.unloc import optimizer_utils
label_desc: Use image node_id 2 for calling the OptimizerUtilsTest obj's underlying member method code with example usage: obj.test_optimizer_with_multi_lrs() without return types