class_docstr: Duplex input Vision Transformer model.
  
  Attributes:
  num_classes: Number of output classes.
  patches: Configuration of the patches extracted in the stem of the model.
  hidden_size: Size of the hidden state of the output of model's stem.
  encoder: Configuration of the encoders used in the model.
  dropout_rate: Dropout rate.
  classifier: type of the classifier layer. Options are 'gap', 'gmp', 'gsp',
    'token'.
  dtype: JAX data type for activations.
class_docstr_tok: ["Duplex","input","Vision","Transformer","model",".","Attributes",":","num_classes",":","Number","of","output","classes",".","patches",":","Configuration","of","the","patches","extracted","in","the","stem","of","the","model",".","hidden_size",":","Size","of","the","hidden","state","of","the","output","of","model","'s","stem",".","encoder",":","Configuration","of","the","encoders","used","in","the","model",".","dropout_rate",":","Dropout","rate",".","classifier",":","type","of","the","classifier","layer",".","Options","are","'gap","'",",","'gmp","'",",","'gsp","'",",","'token","'",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      context_x: jnp.ndarray,
      *,
      train: bool,
      debug: bool = False,
  ):
      # Extracting patches and then embedding is in fact a single convolution.
      fh, fw = self.patches.input_size
      x = nn.Conv(
          self.hidden_size,
          (fh, fw),
          strides=(fh, fw),
          padding="VALID",
          name="embedding",
      )(x)
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])
  
      # Repeat the extraction of patches for the context image.
      context_fh, context_fw = self.patches.context_size
      context_x = nn.Conv(
          self.hidden_size,
          (context_fh, context_fw),
          strides=(context_fh, context_fw),
          padding="VALID",
          name="context_embedding",
      )(context_x)
      context_n, context_h, context_w, context_c = context_x.shape
      context_x = jnp.reshape(
          context_x, [context_n, context_h * context_w, context_c]
      )
  
      # If we want to add a class token, add it to only input (not context).
      if self.classifier == "token":
          cls = self.param(
              "cls", nn.initializers.zeros, (1, 1, c), x.dtype
          )
          cls = jnp.tile(cls, [n, 1, 1])
          x = jnp.concatenate([cls, x], axis=1)
  
      # Add potitional embedding for both input and context
      x = vit.AddPositionEmbs(
          posemb_init=nn.initializers.normal(stddev=0.02),
          name="posembed_input",
      )(x)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      context_x = vit.AddPositionEmbs(
          posemb_init=nn.initializers.normal(stddev=0.02),
          name="posembed_context",
      )(context_x)
      context_x = nn.Dropout(rate=self.dropout_rate)(
          context_x, deterministic=not train
      )
  
      # Optional encoding for input.
      if self.encoder.get("input"):
          x = Encoder(
              mlp_dim=self.encoder.input.mlp_dim,
              num_layers=self.encoder.input.num_layers,
              num_heads=self.encoder.input.num_heads,
              dropout_rate=self.encoder.input.dropout_rate,
              attention_dropout_rate=self.encoder.input.attention_dropout_rate,
              stochastic_depth=self.encoder.input.stochastic_depth,
              dtype=self.dtype,
              name="Input_Transformer",
          )(x, train=train)
  
      # Optional encoding for context.
      if self.encoder.get("context"):
          context_x = Encoder(
              mlp_dim=self.encoder.context.mlp_dim,
              num_layers=self.encoder.context.num_layers,
              num_heads=self.encoder.context.num_heads,
              dropout_rate=self.encoder.context.dropout_rate,
              attention_dropout_rate=self.encoder.context.attention_dropout_rate,
              stochastic_depth=self.encoder.context.stochastic_depth,
              dtype=self.dtype,
              name="Context_Transformer",
          )(context_x, train=train)
  
      # Concat input and context for optionally extra processing on both.
      x = jnp.concatenate([x, context_x], axis=1)
  
      # Optional encoding for input+context.
      if self.encoder.get("fused"):
          x = Encoder(
              mlp_dim=self.encoder.fused.mlp_dim,
              num_layers=self.encoder.fused.num_layers,
              num_heads=self.encoder.fused.num_heads,
              dropout_rate=self.encoder.fused.dropout_rate,
              attention_dropout_rate=self.encoder.fused.attention_dropout_rate,
              stochastic_depth=self.encoder.fused.stochastic_depth,
              dtype=self.dtype,
              name="Transformer",
          )(x, train=train)
  
      if self.classifier in ("token", "0"):
          if (self.encoder.get("fused") is None) and (
              self.encoder.get("context") is not None
          ):
              raise ValueError(
                  "You are encoding the context but"
                  "not using it since the CLS token come from the input"
                  "encoder. Either use gap/gmp/gsp or add a few layers"
                  "of fused encoder."
              )
          x = x[:, 0]
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x = fn(x, axis=1)
  
      x = nn_layers.IdentityLayer(name="pre_logits")(x)
      x = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x)
      return x
  
file_path: scenic/scenic/projects/tasseo/duplex_vit.py
imports: from typing import Any
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the DuplexViT obj's underlying member method code with example usage: obj.__call__(x, context_x) and returns: x