class_docstr: Vision Transformer model for multi-label classification task.
docstr: Loads parameters from an AugReg checkpoint.
  
  See
  https://github.com/google-research/vision_transformer/
  and
  https://arxiv.org/abs/2106.10270
  for more information about these pre-trained models.
  
  Args:
    train_state: A raw TrainState for the model.
    params_path: Path to an Augreg checkpoint. The model config is read from
      the filename (e.g. a B/32 model starts with "B_32-").
    model_cfg: Configuration of the model. Usually used for some asserts.
  
  Returns:
    Updated train_state with params replaced with the ones read from the
    AugReg checkpoint.
docstr_tok: ["Loads","parameters","from","an","AugReg","checkpoint",".","See","https",":","\/\/github.com\/google-research\/vision_transformer\/","and","https",":","\/\/arxiv.org\/abs\/2106.10270","for","more","information","about","these","pre-trained","models",".","Args",":","train_state",":","A","raw","TrainState","for","the","model",".","params_path",":","Path","to","an","Augreg","checkpoint",".","The","model","config","is","read","from","the","filename","(","e.g",".","a","B\/32","model","starts","with","``","B_32-","''",")",".","model_cfg",":","Configuration","of","the","model",".","Usually","used","for","some","asserts",".","Returns",":","Updated","train_state","with","params","replaced","with","the","ones","read","from","the","AugReg","checkpoint","."]
code: def load_augreg_params(
      self,
      train_state: Any,
      params_path: str,
      model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Loads parameters from an AugReg checkpoint.
  
      See
      https://github.com/google-research/vision_transformer/
      and
      https://arxiv.org/abs/2106.10270
      for more information about these pre-trained models.
  
      Args:
        train_state: A raw TrainState for the model.
        params_path: Path to an Augreg checkpoint. The model config is read from
          the filename (e.g. a B/32 model starts with "B_32-").
        model_cfg: Configuration of the model. Usually used for some asserts.
  
      Returns:
        Updated train_state with params replaced with the ones read from the
        AugReg checkpoint.
      """
      restored_model_cfg = _get_augreg_cfg(params_path)
      assert tuple(restored_model_cfg.patches.size) == tuple(
          model_cfg.patches.size
      )
      assert restored_model_cfg.hidden_size == model_cfg.hidden_size
      assert restored_model_cfg.mlp_dim == model_cfg.mlp_dim
      assert restored_model_cfg.num_layers == model_cfg.num_layers
      assert restored_model_cfg.num_heads == model_cfg.num_heads
      assert restored_model_cfg.classifier == model_cfg.classifier
  
      flattened = np.load(gfile.GFile(params_path, "rb"))
      restored_params = flax.traverse_util.unflatten_dict(
          {tuple(k.split("/")): v for k, v in flattened.items()}
      )
      restored_params["output_projection"] = restored_params.pop("head")
  
      if "optimizer" in train_state:
          # TODO(dehghani): Remove support for flax optim.
          params = flax.core.unfreeze(train_state.optimizer.target)
          _merge_params(
              params, restored_params, model_cfg, restored_model_cfg
          )
          return train_state.replace(
              optimizer=train_state.optimizer.replace(
                  target=flax.core.freeze(params)
              )
          )
      else:
          params = flax.core.unfreeze(train_state.params)
          _merge_params(
              params, restored_params, model_cfg, restored_model_cfg
          )
          return train_state.replace(params=flax.core.freeze(params))
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
  from tensorflow.io import gfile
label_desc: Use image node_id 4 for calling the ViTMultiLabelClassificationModel obj's underlying member method code with example usage: obj.load_augreg_params(train_state, params_path, model_cfg) and returns: train_state, train_state