CLIP (Contrastive Language–Image Pre-training)

CLIP (Contrastive Language–Image Pre-training) is a neural network architecture that learns visual concepts from natural language supervision. It is trained on a large dataset of image-text pairs to create a unified vision-language model that can understand both images and text in a shared semantic space.

CLIP consists of two main components: 1. A vision encoder (Vision Transformer) that processes images into visual features 2. A text encoder (Transformer) that processes text into textual features

The model is trained using contrastive learning, where it learns to maximize the cosine similarity between the embeddings of matching image-text pairs while minimizing it for non-matching pairs. This allows CLIP to perform zero-shot classification by comparing image embeddings with text embeddings of potential labels.

CLIP was introduced in the paper "Learning Transferable Visual Models From Natural Language Supervision" and has shown remarkable zero-shot generalization capabilities across a wide range of visual classification tasks. The CLIP model combines a Vision Transformer and a Text Transformer to learn joint representations of images and text. It is trained to maximize the similarity between matching image-text pairs while minimizing similarity between non-matching pairs.

jimm.models.clip.CLIPVisionModel

Bases: Module

Source code in src/jimm/models/clip/clip_model.py

class CLIPVisionModel(nnx.Module):
    def __init__(
        self,
        image_resolution: int,
        vision_layers: int,
        vision_width: int,
        vision_patch_size: int,
        transformer_width: int,
        use_gradient_checkpointing: bool = False,
        rngs: rnglib.Rngs = nnx.Rngs(0),
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        mesh: Mesh | None = None,
        mesh_rules: MeshRules = DEFAULT_SHARDING,
    ):
        """Initialize the Vision Encoder with projection.

        Args:
            image_resolution (int): The resolution of the input images.
            vision_layers (int): The number of layers in the vision transformer.
            vision_width (int): The width of the vision transformer.
            vision_patch_size (int): The patch size of the vision transformer.
            transformer_width (int): The output dimension after projection.
            use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
            rngs (rnglib.Rngs, optional): The random number generator state. Defaults to nnx.Rngs(0).
            dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
            mesh (Mesh | None, optional): The device mesh for parameter sharding. Defaults to None.
            mesh_rules (MeshRules, optional): Logical axis sharding rules. Defaults to DEFAULT_SHARDING.
        """
        self.vision_layers = vision_layers
        self.vision_width = vision_width
        self.vision_patch_size = vision_patch_size
        self.transformer_width = transformer_width
        self.dtype = dtype

        vision_heads = vision_width // 64

        self.encoder = VisionTransformerBase(
            img_size=image_resolution,
            patch_size=vision_patch_size,
            in_channels=3,
            hidden_size=vision_width,
            num_layers=vision_layers,
            num_heads=vision_heads,
            mlp_dim=vision_width * 4,
            use_pre_norm=True,
            use_patch_bias=False,
            use_quick_gelu=True,
            use_gradient_checkpointing=use_gradient_checkpointing,
            pooling_type="CLS",
            layernorm_epsilon=1e-5,
            dtype=dtype,
            param_dtype=param_dtype,
            mesh=mesh,
            rngs=rngs,
            mesh_rules=mesh_rules,
        )
        self.visual_projection = nnx.Linear(
            vision_width,
            transformer_width,
            use_bias=False,
            dtype=dtype,
            param_dtype=param_dtype,
            rngs=rngs,
            kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), mesh_rules("visual_proj_in", "visual_proj_out")),
        )

    def __call__(self, image: Float[Array, "batch height width channels"], do_projection: bool = True) -> Float[Array, "batch vision_width_or_transformer_width"]:
        """Encode images into embeddings.

        Args:
            image (Float[Array, "batch height width channels"]): Batch of input images.
            do_projection (bool): Whether to apply the visual projection layer. Defaults to True.

        Returns:
            Float[Array, "batch vision_width_or_transformer_width"]: Image embeddings.
            Shape depends on do_projection: vision_width if False, transformer_width if True.
        """
        features = self.encoder(image)
        if do_projection:
            return self.visual_projection(features)
        return features

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        use_pytorch: bool = False,
        mesh: Mesh | None = None,
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        use_gradient_checkpointing: bool = False,
        rngs: rnglib.Rngs = nnx.Rngs(0),
    ) -> "CLIPVisionModel":
        """Load a pretrained vision encoder from a CLIP checkpoint.

        Args:
            model_name_or_path (str): Path to local weights or HuggingFace model ID.
            use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
            mesh (Mesh | None): Optional device mesh for parameter sharding. Defaults to None.
            dtype (DTypeLike): Data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike): Data type for parameters. Defaults to jnp.float32.
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            rngs (rnglib.Rngs): Random number generator keys. Defaults to nnx.Rngs(0).

        Returns:
            CLIPVisionModel: Pretrained CLIP vision model
        """
        from .params import load_vision_from_pretrained

        return load_vision_from_pretrained(cls, model_name_or_path, use_pytorch, mesh, dtype, param_dtype, use_gradient_checkpointing, rngs)

    def save_pretrained(self, save_directory: str) -> None:
        """Save model weights and config in HuggingFace format.

        Args:
            save_directory (str): Directory path where the model will be saved.
        """
        from .params import save_vision_pretrained

        save_vision_pretrained(self, save_directory)

__call__(image, do_projection=True)

Encode images into embeddings.

Parameters:

Name	Type	Description	Default
image	Float[Array, 'batch height width channels']	Batch of input images.	required
do_projection	bool	Whether to apply the visual projection layer. Defaults to True.	True

Returns:

Type	Description
Float[Array, 'batch vision_width_or_transformer_width']	Float[Array, "batch vision_width_or_transformer_width"]: Image embeddings.
Float[Array, 'batch vision_width_or_transformer_width']	Shape depends on do_projection: vision_width if False, transformer_width if True.

Source code in src/jimm/models/clip/clip_model.py

def __call__(self, image: Float[Array, "batch height width channels"], do_projection: bool = True) -> Float[Array, "batch vision_width_or_transformer_width"]:
    """Encode images into embeddings.

    Args:
        image (Float[Array, "batch height width channels"]): Batch of input images.
        do_projection (bool): Whether to apply the visual projection layer. Defaults to True.

    Returns:
        Float[Array, "batch vision_width_or_transformer_width"]: Image embeddings.
        Shape depends on do_projection: vision_width if False, transformer_width if True.
    """
    features = self.encoder(image)
    if do_projection:
        return self.visual_projection(features)
    return features

__init__(image_resolution, vision_layers, vision_width, vision_patch_size, transformer_width, use_gradient_checkpointing=False, rngs=nnx.Rngs(0), dtype=jnp.float32, param_dtype=jnp.float32, mesh=None, mesh_rules=DEFAULT_SHARDING)

Initialize the Vision Encoder with projection.

Parameters:

Name	Type	Description	Default
image_resolution	int	The resolution of the input images.	required
vision_layers	int	The number of layers in the vision transformer.	required
vision_width	int	The width of the vision transformer.	required
vision_patch_size	int	The patch size of the vision transformer.	required
transformer_width	int	The output dimension after projection.	required
use_gradient_checkpointing	bool	Whether to use gradient checkpointing. Defaults to False.	False
rngs	Rngs	The random number generator state. Defaults to nnx.Rngs(0).	Rngs(0)
dtype	DTypeLike	The data type for computations. Defaults to jnp.float32.	float32
param_dtype	DTypeLike	The data type for parameters. Defaults to jnp.float32.	float32
mesh	Mesh | None	The device mesh for parameter sharding. Defaults to None.	None
mesh_rules	MeshRules	Logical axis sharding rules. Defaults to DEFAULT_SHARDING.	DEFAULT_SHARDING

Source code in src/jimm/models/clip/clip_model.py

def __init__(
    self,
    image_resolution: int,
    vision_layers: int,
    vision_width: int,
    vision_patch_size: int,
    transformer_width: int,
    use_gradient_checkpointing: bool = False,
    rngs: rnglib.Rngs = nnx.Rngs(0),
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    mesh: Mesh | None = None,
    mesh_rules: MeshRules = DEFAULT_SHARDING,
):
    """Initialize the Vision Encoder with projection.

    Args:
        image_resolution (int): The resolution of the input images.
        vision_layers (int): The number of layers in the vision transformer.
        vision_width (int): The width of the vision transformer.
        vision_patch_size (int): The patch size of the vision transformer.
        transformer_width (int): The output dimension after projection.
        use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
        rngs (rnglib.Rngs, optional): The random number generator state. Defaults to nnx.Rngs(0).
        dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
        mesh (Mesh | None, optional): The device mesh for parameter sharding. Defaults to None.
        mesh_rules (MeshRules, optional): Logical axis sharding rules. Defaults to DEFAULT_SHARDING.
    """
    self.vision_layers = vision_layers
    self.vision_width = vision_width
    self.vision_patch_size = vision_patch_size
    self.transformer_width = transformer_width
    self.dtype = dtype

    vision_heads = vision_width // 64

    self.encoder = VisionTransformerBase(
        img_size=image_resolution,
        patch_size=vision_patch_size,
        in_channels=3,
        hidden_size=vision_width,
        num_layers=vision_layers,
        num_heads=vision_heads,
        mlp_dim=vision_width * 4,
        use_pre_norm=True,
        use_patch_bias=False,
        use_quick_gelu=True,
        use_gradient_checkpointing=use_gradient_checkpointing,
        pooling_type="CLS",
        layernorm_epsilon=1e-5,
        dtype=dtype,
        param_dtype=param_dtype,
        mesh=mesh,
        rngs=rngs,
        mesh_rules=mesh_rules,
    )
    self.visual_projection = nnx.Linear(
        vision_width,
        transformer_width,
        use_bias=False,
        dtype=dtype,
        param_dtype=param_dtype,
        rngs=rngs,
        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), mesh_rules("visual_proj_in", "visual_proj_out")),
    )

from_pretrained(model_name_or_path, use_pytorch=False, mesh=None, dtype=jnp.float32, param_dtype=jnp.float32, use_gradient_checkpointing=False, rngs=nnx.Rngs(0)) classmethod

Load a pretrained vision encoder from a CLIP checkpoint.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	Path to local weights or HuggingFace model ID.	required
use_pytorch	bool	Whether to load from PyTorch weights. Defaults to False.	False
mesh	Mesh | None	Optional device mesh for parameter sharding. Defaults to None.	None
dtype	DTypeLike	Data type for computations. Defaults to jnp.float32.	float32
param_dtype	DTypeLike	Data type for parameters. Defaults to jnp.float32.	float32
use_gradient_checkpointing	bool	Whether to use gradient checkpointing. Defaults to False.	False
rngs	Rngs	Random number generator keys. Defaults to nnx.Rngs(0).	Rngs(0)

Returns:

Name	Type	Description
CLIPVisionModel	CLIPVisionModel	Pretrained CLIP vision model

Source code in src/jimm/models/clip/clip_model.py

@classmethod
def from_pretrained(
    cls,
    model_name_or_path: str,
    use_pytorch: bool = False,
    mesh: Mesh | None = None,
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    use_gradient_checkpointing: bool = False,
    rngs: rnglib.Rngs = nnx.Rngs(0),
) -> "CLIPVisionModel":
    """Load a pretrained vision encoder from a CLIP checkpoint.

    Args:
        model_name_or_path (str): Path to local weights or HuggingFace model ID.
        use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
        mesh (Mesh | None): Optional device mesh for parameter sharding. Defaults to None.
        dtype (DTypeLike): Data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike): Data type for parameters. Defaults to jnp.float32.
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        rngs (rnglib.Rngs): Random number generator keys. Defaults to nnx.Rngs(0).

    Returns:
        CLIPVisionModel: Pretrained CLIP vision model
    """
    from .params import load_vision_from_pretrained

    return load_vision_from_pretrained(cls, model_name_or_path, use_pytorch, mesh, dtype, param_dtype, use_gradient_checkpointing, rngs)

save_pretrained(save_directory)

Save model weights and config in HuggingFace format.

Parameters:

Name	Type	Description	Default
save_directory	str	Directory path where the model will be saved.	required

Source code in src/jimm/models/clip/clip_model.py

def save_pretrained(self, save_directory: str) -> None:
    """Save model weights and config in HuggingFace format.

    Args:
        save_directory (str): Directory path where the model will be saved.
    """
    from .params import save_vision_pretrained

    save_vision_pretrained(self, save_directory)

jimm.models.clip.CLIP

Bases: Module

Source code in src/jimm/models/clip/clip_model.py

class CLIP(nnx.Module):
    def __init__(
        self,
        image_resolution: int,
        vision_layers: int,
        vision_width: int,
        vision_patch_size: int,
        context_length: int,
        vocab_size: int,
        transformer_width: int,
        transformer_heads: int,
        transformer_layers: int,
        use_gradient_checkpointing: bool = False,
        rngs: rnglib.Rngs = nnx.Rngs(0),
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        mesh: Mesh | None = None,
        mesh_rules: MeshRules = DEFAULT_SHARDING,
    ):
        """Initialize the CLIP model.

        Args:
            image_resolution (int): The resolution of the input images.
            vision_layers (int): The number of layers in the vision transformer.
            vision_width (int): The width of the vision transformer.
            vision_patch_size (int): The patch size of the vision transformer.
            context_length (int): The length of the context.
            vocab_size (int): The size of the vocabulary.
            transformer_width (int): The width of the transformer.
            transformer_heads (int): The number of attention heads in the transformer.
            transformer_layers (int): The number of layers in the transformer.
            use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
            rngs (rnglib.Rngs, optional): The random number generator state. Defaults to nnx.Rngs(0).
            dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
            mesh (Mesh | None, optional): The device mesh for parameter sharding. Defaults to None.
            mesh_rules (MeshRules, optional): Logical axis sharding rules. Defaults to DEFAULT_SHARDING.
        """
        self.vision_layers = vision_layers
        self.vision_width = vision_width
        self.vision_patch_size = vision_patch_size
        self.context_length = context_length
        self.vocab_size = vocab_size
        self.transformer_width = transformer_width
        self.transformer_heads = transformer_heads
        self.transformer_layers = transformer_layers
        self.dtype = dtype
        self._original_config = None

        self.vision_model = CLIPVisionModel(
            image_resolution=image_resolution,
            vision_layers=vision_layers,
            vision_width=vision_width,
            vision_patch_size=vision_patch_size,
            transformer_width=transformer_width,
            use_gradient_checkpointing=use_gradient_checkpointing,
            rngs=rngs,
            dtype=dtype,
            param_dtype=param_dtype,
            mesh=mesh,
            mesh_rules=mesh_rules,
        )

        self.text_model = CLIPTextModel(
            context_length=context_length,
            vocab_size=vocab_size,
            transformer_width=transformer_width,
            transformer_heads=transformer_heads,
            transformer_layers=transformer_layers,
            use_gradient_checkpointing=use_gradient_checkpointing,
            rngs=rngs,
            dtype=dtype,
            param_dtype=param_dtype,
            mesh=mesh,
            mesh_rules=mesh_rules,
        )

        self.logit_scale = nnx.Param(nnx.with_partitioning(nnx.initializers.ones_init(), ())(rngs.params(), ()))

    def encode_image(self, image: Float[Array, "batch height width channels"], do_projection: bool = True) -> Float[Array, "batch transformer_width"]:
        """Encode images into embeddings.

        Args:
            image (Float[Array, "batch height width channels"]): Batch of input images.
            do_projection (bool): Whether the image encoder should do the visual projection layer. Defaults to true.

        Returns:
            Float[Array, "batch transformer_width"]: Image embeddings.
        """
        return self.vision_model(image, do_projection)

    def encode_text(self, text: Int[Array, "batch context_length"]) -> Float[Array, "batch transformer_width"]:
        """Encode text tokens into embeddings.

        Args:
            text (Int[Array, "batch context_length"]): Batch of token sequences.

        Returns:
            Float[Array, "batch transformer_width"]: Text embeddings.
        """
        return self.text_model(text, do_projection=True)

    def __call__(self, image: Float[Array, "batch height width channels"], text: Int[Array, "batch context_length"]) -> Float[Array, "batch batch"]:
        """Calculate similarity between image and text embeddings.

        Args:
            image (Float[Array, "batch height width channels"]): Batch of input images.
            text (Int[Array, "batch context_length"]): Batch of token sequences.

        Returns:
            Float[Array, "batch batch"]: Similarity scores between all pairs of images and texts.
        """
        image_features: Float[Array, "batch transformer_width"] = self.encode_image(image, do_projection=True)
        text_features: Float[Array, "batch transformer_width"] = self.encode_text(text)

        image_features: Float[Array, "batch transformer_width"] = image_features / jnp.linalg.norm(image_features, axis=-1, keepdims=True)
        text_features: Float[Array, "batch transformer_width"] = text_features / jnp.linalg.norm(text_features, axis=-1, keepdims=True)

        logit_scale: Float[Array, ""] = jnp.exp(self.logit_scale)
        logits: Float[Array, "batch batch"] = logit_scale * image_features @ text_features.T
        return logits

    @classmethod
    def from_pretrained(
        cls,
        model_name_or_path: str,
        use_pytorch: bool = False,
        mesh: Mesh | None = None,
        dtype: DTypeLike = jnp.float32,
        param_dtype: DTypeLike = jnp.float32,
        use_gradient_checkpointing: bool = False,
        rngs: rnglib.Rngs = nnx.Rngs(0),
    ) -> "CLIP":
        """Load a pretrained CLIP model from a local path or HuggingFace Hub.

        Args:
            model_name_or_path (str): Path to local weights or HuggingFace model ID.
            use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
            mesh (Mesh | None): Optional device mesh for parameter sharding. Defaults to None.
            dtype (DTypeLike): Data type for computations. Defaults to jnp.float32.
            param_dtype (DTypeLike): Data type for parameters. Defaults to jnp.float32.
            use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
            rngs (rnglib.Rngs): Random number generator keys. Defaults to nnx.Rngs(0).

        Returns:
            CLIP: Pretrained CLIP model
        """
        from .params import load_from_pretrained

        return load_from_pretrained(cls, model_name_or_path, use_pytorch, mesh, dtype, param_dtype, use_gradient_checkpointing, rngs)

    def save_pretrained(self, save_directory: str) -> None:
        """Save the model weights and config in HuggingFace format.

        Args:
            save_directory (str): Directory path where the model will be saved.
        """
        from .params import save_pretrained

        save_pretrained(self, save_directory)

__call__(image, text)

Calculate similarity between image and text embeddings.

Parameters:

Name	Type	Description	Default
image	Float[Array, 'batch height width channels']	Batch of input images.	required
text	Int[Array, 'batch context_length']	Batch of token sequences.	required

Returns:

Type	Description
Float[Array, 'batch batch']	Float[Array, "batch batch"]: Similarity scores between all pairs of images and texts.

Source code in src/jimm/models/clip/clip_model.py

def __call__(self, image: Float[Array, "batch height width channels"], text: Int[Array, "batch context_length"]) -> Float[Array, "batch batch"]:
    """Calculate similarity between image and text embeddings.

    Args:
        image (Float[Array, "batch height width channels"]): Batch of input images.
        text (Int[Array, "batch context_length"]): Batch of token sequences.

    Returns:
        Float[Array, "batch batch"]: Similarity scores between all pairs of images and texts.
    """
    image_features: Float[Array, "batch transformer_width"] = self.encode_image(image, do_projection=True)
    text_features: Float[Array, "batch transformer_width"] = self.encode_text(text)

    image_features: Float[Array, "batch transformer_width"] = image_features / jnp.linalg.norm(image_features, axis=-1, keepdims=True)
    text_features: Float[Array, "batch transformer_width"] = text_features / jnp.linalg.norm(text_features, axis=-1, keepdims=True)

    logit_scale: Float[Array, ""] = jnp.exp(self.logit_scale)
    logits: Float[Array, "batch batch"] = logit_scale * image_features @ text_features.T
    return logits

__init__(image_resolution, vision_layers, vision_width, vision_patch_size, context_length, vocab_size, transformer_width, transformer_heads, transformer_layers, use_gradient_checkpointing=False, rngs=nnx.Rngs(0), dtype=jnp.float32, param_dtype=jnp.float32, mesh=None, mesh_rules=DEFAULT_SHARDING)

Initialize the CLIP model.

Parameters:

Name	Type	Description	Default
image_resolution	int	The resolution of the input images.	required
vision_layers	int	The number of layers in the vision transformer.	required
vision_width	int	The width of the vision transformer.	required
vision_patch_size	int	The patch size of the vision transformer.	required
context_length	int	The length of the context.	required
vocab_size	int	The size of the vocabulary.	required
transformer_width	int	The width of the transformer.	required
transformer_heads	int	The number of attention heads in the transformer.	required
transformer_layers	int	The number of layers in the transformer.	required
use_gradient_checkpointing	bool	Whether to use gradient checkpointing. Defaults to False.	False
rngs	Rngs	The random number generator state. Defaults to nnx.Rngs(0).	Rngs(0)
dtype	DTypeLike	The data type for computations. Defaults to jnp.float32.	float32
param_dtype	DTypeLike	The data type for parameters. Defaults to jnp.float32.	float32
mesh	Mesh | None	The device mesh for parameter sharding. Defaults to None.	None
mesh_rules	MeshRules	Logical axis sharding rules. Defaults to DEFAULT_SHARDING.	DEFAULT_SHARDING

Source code in src/jimm/models/clip/clip_model.py

def __init__(
    self,
    image_resolution: int,
    vision_layers: int,
    vision_width: int,
    vision_patch_size: int,
    context_length: int,
    vocab_size: int,
    transformer_width: int,
    transformer_heads: int,
    transformer_layers: int,
    use_gradient_checkpointing: bool = False,
    rngs: rnglib.Rngs = nnx.Rngs(0),
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    mesh: Mesh | None = None,
    mesh_rules: MeshRules = DEFAULT_SHARDING,
):
    """Initialize the CLIP model.

    Args:
        image_resolution (int): The resolution of the input images.
        vision_layers (int): The number of layers in the vision transformer.
        vision_width (int): The width of the vision transformer.
        vision_patch_size (int): The patch size of the vision transformer.
        context_length (int): The length of the context.
        vocab_size (int): The size of the vocabulary.
        transformer_width (int): The width of the transformer.
        transformer_heads (int): The number of attention heads in the transformer.
        transformer_layers (int): The number of layers in the transformer.
        use_gradient_checkpointing (bool, optional): Whether to use gradient checkpointing. Defaults to False.
        rngs (rnglib.Rngs, optional): The random number generator state. Defaults to nnx.Rngs(0).
        dtype (DTypeLike, optional): The data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike, optional): The data type for parameters. Defaults to jnp.float32.
        mesh (Mesh | None, optional): The device mesh for parameter sharding. Defaults to None.
        mesh_rules (MeshRules, optional): Logical axis sharding rules. Defaults to DEFAULT_SHARDING.
    """
    self.vision_layers = vision_layers
    self.vision_width = vision_width
    self.vision_patch_size = vision_patch_size
    self.context_length = context_length
    self.vocab_size = vocab_size
    self.transformer_width = transformer_width
    self.transformer_heads = transformer_heads
    self.transformer_layers = transformer_layers
    self.dtype = dtype
    self._original_config = None

    self.vision_model = CLIPVisionModel(
        image_resolution=image_resolution,
        vision_layers=vision_layers,
        vision_width=vision_width,
        vision_patch_size=vision_patch_size,
        transformer_width=transformer_width,
        use_gradient_checkpointing=use_gradient_checkpointing,
        rngs=rngs,
        dtype=dtype,
        param_dtype=param_dtype,
        mesh=mesh,
        mesh_rules=mesh_rules,
    )

    self.text_model = CLIPTextModel(
        context_length=context_length,
        vocab_size=vocab_size,
        transformer_width=transformer_width,
        transformer_heads=transformer_heads,
        transformer_layers=transformer_layers,
        use_gradient_checkpointing=use_gradient_checkpointing,
        rngs=rngs,
        dtype=dtype,
        param_dtype=param_dtype,
        mesh=mesh,
        mesh_rules=mesh_rules,
    )

    self.logit_scale = nnx.Param(nnx.with_partitioning(nnx.initializers.ones_init(), ())(rngs.params(), ()))

encode_image(image, do_projection=True)

Encode images into embeddings.

Parameters:

Name	Type	Description	Default
image	Float[Array, 'batch height width channels']	Batch of input images.	required
do_projection	bool	Whether the image encoder should do the visual projection layer. Defaults to true.	True

Returns:

Type	Description
Float[Array, 'batch transformer_width']	Float[Array, "batch transformer_width"]: Image embeddings.

Source code in src/jimm/models/clip/clip_model.py

def encode_image(self, image: Float[Array, "batch height width channels"], do_projection: bool = True) -> Float[Array, "batch transformer_width"]:
    """Encode images into embeddings.

    Args:
        image (Float[Array, "batch height width channels"]): Batch of input images.
        do_projection (bool): Whether the image encoder should do the visual projection layer. Defaults to true.

    Returns:
        Float[Array, "batch transformer_width"]: Image embeddings.
    """
    return self.vision_model(image, do_projection)

encode_text(text)

Encode text tokens into embeddings.

Parameters:

Name	Type	Description	Default
text	Int[Array, 'batch context_length']	Batch of token sequences.	required

Returns:

Type	Description
Float[Array, 'batch transformer_width']	Float[Array, "batch transformer_width"]: Text embeddings.

Source code in src/jimm/models/clip/clip_model.py

def encode_text(self, text: Int[Array, "batch context_length"]) -> Float[Array, "batch transformer_width"]:
    """Encode text tokens into embeddings.

    Args:
        text (Int[Array, "batch context_length"]): Batch of token sequences.

    Returns:
        Float[Array, "batch transformer_width"]: Text embeddings.
    """
    return self.text_model(text, do_projection=True)

from_pretrained(model_name_or_path, use_pytorch=False, mesh=None, dtype=jnp.float32, param_dtype=jnp.float32, use_gradient_checkpointing=False, rngs=nnx.Rngs(0)) classmethod

Load a pretrained CLIP model from a local path or HuggingFace Hub.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	Path to local weights or HuggingFace model ID.	required
use_pytorch	bool	Whether to load from PyTorch weights. Defaults to False.	False
mesh	Mesh | None	Optional device mesh for parameter sharding. Defaults to None.	None
dtype	DTypeLike	Data type for computations. Defaults to jnp.float32.	float32
param_dtype	DTypeLike	Data type for parameters. Defaults to jnp.float32.	float32
use_gradient_checkpointing	bool	Whether to use gradient checkpointing. Defaults to False.	False
rngs	Rngs	Random number generator keys. Defaults to nnx.Rngs(0).	Rngs(0)

Returns:

Name	Type	Description
CLIP	CLIP	Pretrained CLIP model

Source code in src/jimm/models/clip/clip_model.py

@classmethod
def from_pretrained(
    cls,
    model_name_or_path: str,
    use_pytorch: bool = False,
    mesh: Mesh | None = None,
    dtype: DTypeLike = jnp.float32,
    param_dtype: DTypeLike = jnp.float32,
    use_gradient_checkpointing: bool = False,
    rngs: rnglib.Rngs = nnx.Rngs(0),
) -> "CLIP":
    """Load a pretrained CLIP model from a local path or HuggingFace Hub.

    Args:
        model_name_or_path (str): Path to local weights or HuggingFace model ID.
        use_pytorch (bool): Whether to load from PyTorch weights. Defaults to False.
        mesh (Mesh | None): Optional device mesh for parameter sharding. Defaults to None.
        dtype (DTypeLike): Data type for computations. Defaults to jnp.float32.
        param_dtype (DTypeLike): Data type for parameters. Defaults to jnp.float32.
        use_gradient_checkpointing (bool): Whether to use gradient checkpointing. Defaults to False.
        rngs (rnglib.Rngs): Random number generator keys. Defaults to nnx.Rngs(0).

    Returns:
        CLIP: Pretrained CLIP model
    """
    from .params import load_from_pretrained

    return load_from_pretrained(cls, model_name_or_path, use_pytorch, mesh, dtype, param_dtype, use_gradient_checkpointing, rngs)

save_pretrained(save_directory)

Save the model weights and config in HuggingFace format.

Parameters:

Name	Type	Description	Default
save_directory	str	Directory path where the model will be saved.	required

Source code in src/jimm/models/clip/clip_model.py

def save_pretrained(self, save_directory: str) -> None:
    """Save the model weights and config in HuggingFace format.

    Args:
        save_directory (str): Directory path where the model will be saved.
    """
    from .params import save_pretrained

    save_pretrained(self, save_directory)