class_docstr: ViT model.
class_docstr_tok: ["ViT","model","."]
code: def __call__(self, image, *, train=False):
      out = {}
  
      # Patch extraction
      x = out["stem"] = nn.Conv(
          self.width,
          self.patch_size,
          strides=self.patch_size,
          padding="VALID",
          name="embedding",
      )(image)
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])
  
      # Add posemb before adding extra token.
      x = out["with_posemb"] = x + bv_vit.get_posemb(
          self, self.posemb, (h, w), c, "pos_embedding", x.dtype
      )
  
      if self.pool_type == "tok":
          cls = self.param(
              "cls", nn.initializers.zeros, (1, 1, c), x.dtype
          )
          x = jnp.concatenate([jnp.tile(cls, [n, 1, 1]), x], axis=1)
  
      n, l, c = x.shape  # pylint: disable=unused-variable
      x = nn.Dropout(rate=self.dropout)(x, not train)
      x, out["encoder"] = Encoder(
          depth=self.depth,
          mlp_dim=self.mlp_dim,
          num_heads=self.num_heads,
          dropout=self.dropout,
          dtype=self.dtype,
          num_frozen_layers=self.num_frozen_layers,
          name="Transformer",
      )(x, deterministic=not train)
  
      encoded = out["encoded"] = x
  
      if self.pool_type == "map":
          x = out["head_input"] = bv_vit.MAPHead(
              num_heads=self.num_heads, mlp_dim=self.mlp_dim
          )(x)
      elif self.pool_type == "gap":
          x = out["head_input"] = jnp.mean(x, axis=1)
      elif self.pool_type == "0":
          x = out["head_input"] = x[:, 0]
      elif self.pool_type == "tok":
          x = out["head_input"] = x[:, 0]
          encoded = encoded[:, 1:]
      else:
          raise ValueError(f"Unknown pool type: '{self.pool_type}'")
  
      x_2d = jnp.reshape(encoded, [n, h, w, -1])
  
      if self.rep_size:
          rep_size = (
              self.width if self.rep_size is True else self.rep_size
          )  # pylint: disable=g-bool-id-comparison
          hid = nn.Dense(rep_size, name="pre_logits", dtype=self.dtype)
          # NOTE: In the past we did not include tanh in pre_logits.
          # For few-shot, it should not matter much, as it whitens anyways.
          x_2d = nn.tanh(hid(x_2d))
          x = nn.tanh(hid(x))
  
      out["pre_logits_2d"] = x_2d
      out["pre_logits"] = x
  
      if self.num_classes:
          kw = (
              {"kernel_init": nn.initializers.zeros}
              if self.head_zeroinit
              else {}
          )
          head = nn.Dense(
              self.num_classes, name="head", dtype=self.dtype, **kw
          )
          x_2d = out["logits_2d"] = head(x_2d)
          x = out["logits"] = head(x)
      x = jnp.array(x, dtype=self.dtype)
      return x, out
  
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 the _Model obj's underlying member method code with example usage: obj.__call__(image) and returns: x,  out