code: def Model(
      num_classes,
      num_frozen_layers=-1,
      dtype=jnp.bfloat16,
      variant=None,
      **kw,
  ):  # pylint: disable=invalid-name
      """Factory function, because linen really don't like what I'm doing!"""
      vit_config = bv_vit.decode_variant(variant)
      if isinstance(num_frozen_layers, float):
          num_frozen_layers = int(
              num_frozen_layers * vit_config["depth"]
          )
      return _Model(
          num_classes=num_classes,
          num_frozen_layers=num_frozen_layers,
          dtype=dtype,
          **{**vit_config, **kw},
      )
  
file_path: scenic/scenic/projects/knowledge_visual_language/models/vit.py
imports: from typing import Optional, Sequence, Union
  from big_vision.models import vit
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: Model(num_classes, num_frozen_layers, dtype, variant) and returns: _Model