class_docstr: PolyVit.
class_docstr_tok: ["PolyVit","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      targets: Optional[jnp.ndarray] = None,
      *,
      task: Optional[str] = None,
      modality: Optional[str] = "image",
      dataset: Optional[str] = None,
      train: bool,
      debug: bool = False,
  ):
      if self.stochastic_droplayer_config is None:
          stochastic_droplayer_rate = None
      else:
          stochastic_droplayer_rate = (
              self.stochastic_droplayer_config.get(dataset)
          )
  
      x = layers.Tokenizer(
          config=self.modalities_config,
          dtype=self.dtype,
          name="tokenizer",
      )(
          x,
          train=train,
          modality=modality,
          stochastic_droplayer_rate=stochastic_droplayer_rate,
          dataset=dataset,
      )
  
      x = layers.PolyViTEncoder(
          mlp_dim=self.encoder_config.mlp_dim,
          num_tokenizer_layers=self.encoder_config.num_tokenizer_layers,
          num_layers=self.encoder_config.num_layers,
          num_heads=self.encoder_config.num_heads,
          dropout_rate=self.encoder_config.dropout_rate,
          attention_dropout_rate=self.encoder_config.attention_dropout_rate,
          dtype=self.dtype,
          name="vit_encoder",
      )(
          x,
          train=train,
          stochastic_droplayer_rate=stochastic_droplayer_rate,
          dataset=dataset,
      )
  
      if self.encoder_config.get("freeze_body", False):
          x = jax.lax.stop_gradient(x)
  
      elif task in [
          polyvit_base_model.Task.LABEL,
          polyvit_base_model.Task.MULTILABEL,
          polyvit_base_model.Task.MULTIHEADLABEL,
      ]:
          head_config = self.heads_config.label.get(dataset)
  
          x = layers.ClassificationHead(
              num_outputs=head_config.num_classes,
              hid_sizes=head_config.hid_sizes,
              output_proj_zero_init=head_config.get(
                  "output_proj_zero_init", False
              ),
              classifier=head_config.classifier,
              dtype=self.dtype,
              name="classification_head_" + dataset,
          )(x, train=train)
  
          return x
  
      elif task == polyvit_base_model.Task.BOW:
          head_config = self.heads_config.bow.get(dataset)
  
          x = layers.ClassificationHead(
              num_outputs=head_config.vocab_size,
              hid_sizes=head_config.hid_sizes,
              output_proj_zero_init=head_config.get(
                  "output_proj_zero_init", False
              ),
              classifier=head_config.classifier,
              dtype=self.dtype,
              name="bag_of_words_head_" + dataset,
          )(x, train=train)
  
          return x
  
      elif task == polyvit_base_model.Task.FEWSHOT:
          x = layers.FewshotHead(
              self.heads_config.fewshot.pooling_type,
              name="fewshot_head",
          )(x, train=train)
  
          return x
  
      else:
          raise NotImplementedError(
              f"Task {task} is not supported yet."
          )
  
file_path: scenic/scenic/projects/polyvit/model.py
imports: from typing import Any, Callable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.polyvit import layers
  from scenic.projects.polyvit import model_utils
  from scenic.projects.polyvit import polyvit_base_model
label_desc: Use image node_id 1 for calling the PolyVit obj's underlying member method code with example usage: obj.__call__(x, targets) and returns: x, x, x