class_docstr: Implements the Bit ResNet model for classification.
class_docstr_tok: ["Implements","the","Bit","ResNet","model","for","classification","."]
docstr: Updates the train_state with data from `restored_train_state`.
  
  This function is writen to be used for 'fine-tuning' experiments.
  As defined in the ResNet definition above output head is called
  `output_projection` and not loaded since often target tasks have a
  new output head, possibly with different shape.
  
  Args:
    train_state: A raw TrainState for the model.
    restored_train_state: A TrainState that is loaded with parameters/state of
      a pretrained model.
    restored_model_cfg: Configuration of the model from which the
      `restored_train_state` come from. Usually used for some asserts.
  
  Returns:
    Updated train_state.
docstr_tok: ["Updates","the","train_state","with","data","from","`","restored_train_state","`",".","This","function","is","writen","to","be","used","for","'fine-tuning","'","experiments",".","As","defined","in","the","ResNet","definition","above","output","head","is","called","`","output_projection","`","and","not","loaded","since","often","target","tasks","have","a","new","output","head",",","possibly","with","different","shape",".","Args",":","train_state",":","A","raw","TrainState","for","the","model",".","restored_train_state",":","A","TrainState","that","is","loaded","with","parameters\/state","of","a","pretrained","model",".","restored_model_cfg",":","Configuration","of","the","model","from","which","the","`","restored_train_state","`","come","from",".","Usually","used","for","some","asserts",".","Returns",":","Updated","train_state","."]
code: def init_from_train_state(
      self,
      train_state: Any,
      restored_train_state: Any,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the train_state with data from `restored_train_state`.
  
      This function is writen to be used for 'fine-tuning' experiments.
      As defined in the ResNet definition above output head is called
      `output_projection` and not loaded since often target tasks have a
      new output head, possibly with different shape.
  
      Args:
        train_state: A raw TrainState for the model.
        restored_train_state: A TrainState that is loaded with parameters/state of
          a pretrained model.
        restored_model_cfg: Configuration of the model from which the
          `restored_train_state` come from. Usually used for some asserts.
  
      Returns:
        Updated train_state.
      """
      del restored_model_cfg
      if hasattr(train_state, "optimizer"):
          # TODO(dehghani): Remove support for flax optim.
          params = flax.core.unfreeze(train_state.optimizer.target)
          restored_params = flax.core.unfreeze(
              restored_train_state.optimizer.target
          )
      else:
          params = flax.core.unfreeze(train_state.params)
          restored_params = flax.core.unfreeze(
              restored_train_state.params
          )
      # Check all parameters are loaded.
      params_to_load = set(params.keys())
      params_to_load.remove("output_projection")
      for pname, pvalue in restored_params.items():
          if pname == "output_projection":
              # The `output_projection` is used as the name of the linear lyaer at the
              # head of the model that maps the representation to the label space.
              # By default, for finetuning to another dataset, we drop this layer as
              # the label space is different.
              continue
          else:
              if pname not in params:
                  raise ValueError(
                      f"Loaded parameter {pname} doesnt exist in params."
                  )
              params[pname] = pvalue
              params_to_load.remove(pname)
      if params_to_load:
          raise ValueError(
              f"Paramater groups that are not loaded: {params_to_load}"
          )
      logging.info(
          "Parameter summary after initialising from train state:"
      )
      debug_utils.log_param_shapes(params)
      if hasattr(train_state, "optimizer"):
          # TODO(dehghani): Remove support for flax optim.
          return train_state.replace(
              optimizer=train_state.optimizer.replace(
                  target=flax.core.freeze(params)
              ),
              model_state=restored_train_state.model_state,
          )
      else:
          return train_state.replace(
              params=flax.core.freeze(params),
              model_state=restored_train_state.model_state,
          )
  
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 3 for calling the BitResNetClassificationModel obj's underlying member method code with example usage: obj.init_from_train_state(train_state, restored_train_state, restored_model_cfg) and returns: train_state, train_state