class_docstr: Unified Tokenizer class.
  
  Attributes:
    config: Tokenizer config.
    dtype: JAX data type for activations.
class_docstr_tok: ["Unified","Tokenizer","class",".","Attributes",":","config",":","Tokenizer","config",".","dtype",":","JAX","data","type","for","activations","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      *,
      train: bool,
      dataset: str,
      modality: Optional[str],
      stochastic_droplayer_rate: Optional[float],
  ) -> jnp.ndarray:
      if modality == polyvit_base_model.Modality.IMAGE:
          return Tokenizer2D(
              self.config.image.patches,
              self.config.hidden_size,
              self.config.mlp_dim,
              self.config.num_layers,
              self.config.num_heads,
              dropout_rate=self.config.dropout_rate,
              attention_dropout_rate=self.config.attention_dropout_rate,
              add_cls_token=self.config.add_cls_token,
              dtype=self.dtype,
              name="tokenizer2d",
          )(
              x,
              train=train,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dataset=dataset,
          )
      elif modality == polyvit_base_model.Modality.VIDEO:
          return Tokenizer3D(
              self.config.video.patches,
              self.config.hidden_size,
              self.config.video.kernel_init_method,
              self.config.mlp_dim,
              self.config.num_layers,
              self.config.num_heads,
              dropout_rate=self.config.dropout_rate,
              attention_dropout_rate=self.config.attention_dropout_rate,
              add_cls_token=self.config.add_cls_token,
              dtype=self.dtype,
              name="tokenizer3d",
          )(
              x,
              train=train,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dataset=dataset,
          )
      elif modality == polyvit_base_model.Modality.AUDIO:
          return Tokenizer2D(
              self.config.audio.patches,
              self.config.hidden_size,
              self.config.mlp_dim,
              self.config.num_layers,
              self.config.num_heads,
              dropout_rate=self.config.dropout_rate,
              attention_dropout_rate=self.config.attention_dropout_rate,
              add_cls_token=self.config.add_cls_token,
              dtype=self.dtype,
              name="tokenizer_spec",
          )(
              x,
              train=train,
              stochastic_droplayer_rate=stochastic_droplayer_rate,
              dataset=dataset,
          )
      else:
          raise NotImplementedError(
              f"Modality {modality} is not supported yet."
          )
  
imports: import functools
  from typing import Any, Callable, Optional, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.polyvit import polyvit_base_model
  from scenic.projects.vivit import model
label_desc: Use image node_id 1 for calling the Tokenizer obj's underlying member method code with example usage: obj.__call__(x) and returns: Tokenizer2D, Tokenizer3D, Tokenizer2D