class_docstr: Clip model consisting of a vision and text transformer.
  
  Attributes:
    vocab_size: Size of the vocabulary.
    embed_dim: Size of the text and vision embeddings.
    text_features: Number of features in text transformer.
    text_num_layers: Number of text transformer blocks (self-attn + MLP).
    text_num_heads: Number of heads in text transformer.
    vision_features: Number of features in vision transformer.
    vision_num_layers: Number of vision transformer blocks (self-attn + MLP).
    vision_patch_size: Size of patches to embed in vision transformer.
class_docstr_tok: ["Clip","model","consisting","of","a","vision","and","text","transformer",".","Attributes",":","vocab_size",":","Size","of","the","vocabulary",".","embed_dim",":","Size","of","the","text","and","vision","embeddings",".","text_features",":","Number","of","features","in","text","transformer",".","text_num_layers",":","Number","of","text","transformer","blocks","(","self-attn","+","MLP",")",".","text_num_heads",":","Number","of","heads","in","text","transformer",".","vision_features",":","Number","of","features","in","vision","transformer",".","vision_num_layers",":","Number","of","vision","transformer","blocks","(","self-attn","+","MLP",")",".","vision_patch_size",":","Size","of","patches","to","embed","in","vision","transformer","."]
code: def setup(self):
      if isinstance(self.vision_num_layers, (tuple, list)):
          self.vision_num_heads = self.vision_features * 32 // 64
          if self.vision_stochastic_droplayer_rate > 0.0:
              raise ValueError(
                  "ResNet backbone does not support stochastic depth."
              )
          self.visual = ModifiedResNet(
              num_layers=self.vision_num_layers,
              features=self.vision_features,
              num_heads=self.vision_num_heads,
              out_features=None
              if self.vision_return_map
              else self.embed_dim,
          )
      else:
          self.vision_num_heads = self.vision_features // 64
          self.visual = VisionTransformer(
              patch_size=self.vision_patch_size,
              features=self.vision_features,
              num_layers=self.vision_num_layers,
              num_heads=self.vision_num_heads,
              out_features=None
              if self.vision_return_map
              else self.embed_dim,
              stochastic_droplayer_rate=self.vision_stochastic_droplayer_rate,
          )
      self.text = TextEncoder(
          out_features=self.embed_dim,
          vocab_size=self.vocab_size,
          features=self.text_features,
          num_layers=self.text_num_layers,
          num_heads=self.text_num_heads,
          stochastic_droplayer_rate=self.text_stochastic_droplayer_rate,
      )
      self.logit_scale = self.param(
          "logit_scale", jax.nn.initializers.zeros, ()
      )
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the CLIP obj's underlying member method code with example usage: obj.setup() without return types