def plot_learning_rate(config: ml_collections.ConfigDict()):
  """Plots the learning rate in the config.

  Example:
    config1 = vqgan3d_videocc_ft_config.get_config('L')
    config2 = vqgan3d_ucf101_config.get_config()

    colab_utils.plot_learning_rate(config1)

    colab_utils.plot_learning_rate(config2)
    lr_configs = config_utils.get_lr_rsqrt_configs(
        0,
        config1,
        timescale=1_000,
        warmup_steps=0,
        cooldown_steps=15_000,
    )
    colab_utils.plot_learning_rate(lr_configs)

  Args:
    config: a config containing the lr_configs or config is the lr_configs.
  """
  if 'lr_configs' not in config:
    # config is lr_configs
    assert 'factors' in config, 'config does not have a lr_configs field.'
    config = ml_collections.ConfigDict(dict(lr_configs=config))
  learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
  if 'total_steps' in config.lr_configs:
    total_steps = config.lr_configs.total_steps
  elif 'steps_per_cycle' in config.lr_configs:
    total_steps = config.lr_configs.steps_per_cycle    # consine lr
  else:
    raise NotImplementedError('missing total_steps in the learning rate.')

  xs = np.arange(1, total_steps, 100)
  plt.plot(xs, learning_rate_fn(xs))
Plots the learning rate in the config.

  Example:
    config1 = vqgan3d_videocc_ft_config.get_config('L')
    config2 = vqgan3d_ucf101_config.get_config()

    colab_utils.plot_learning_rate(config1)

    colab_utils.plot_learning_rate(config2)
    lr_configs = config_utils.get_lr_rsqrt_configs(
        0,
        config1,
        timescale=1_000,
        warmup_steps=0,
        cooldown_steps=15_000,
    )
    colab_utils.plot_learning_rate(lr_configs)

  Args:
    config: a config containing the lr_configs or config is the lr_configs.
  
videogvt/interfaces/colab_utils.py