Contributing a New Layer

pyc_logo PyC layers are the semantic primitives of the Low-Level API. Every layer extends BaseConceptLayer — a thin pytorch_logo PyTorch nn.Module wrapper that standardises how inputs and outputs are described in terms of concepts, embeddings, and their cardinalities. New layers slot directly into both the Low-Level API and the mid-level ParametricCPD system.

Expand each block below for a step-by-step walkthrough.

Layer Interface

All layers extend BaseConceptLayer. The constructor accepts three optional dimension descriptors — pass an Annotations object instead of an int to make the layer semantics-aware:

class BaseConceptLayer(ABC, torch.nn.Module):
    def __init__(
        self,
        out_concepts: Union[int, Annotations],
        in_concepts:  Union[int, Annotations] = None,
        in_embeddings: Union[int, Annotations] = None,
    ): ...

After super().__init__(...) three resolved-integer attributes are available:

  • self.in_concepts_shape — int (or None if not passed)

  • self.in_embeddings_shape — int (or None if not passed)

  • self.out_concepts_shape — int

These always hold plain integers regardless of whether you passed an Annotations object or an int, so you can use them directly in nn.Linear / nn.Conv.

Naming convention. Layer class names follow <OperationType><InputType>To<OutputType>, e.g.:

  • LinearEmbeddingToConcept — linear map from embeddings to concept scores

  • LinearConceptToConcept — linear map from concept scores to concept scores

  • MixConceptEmbeddingToConcept — mixes concept scores with their embeddings

forward keyword arguments mirror the input types: use concepts= for concept-score tensors and embeddings= for embedding tensors. This lets callers be explicit and lets ParametricCPD route inputs correctly.

Minimal Example

Below is a complete new layer — a concept-score normaliser that applies layer-norm before passing concepts through a linear map:

# torch_concepts/nn/modules/low/predictors/normed.py
from typing import Union
import torch
import torch.nn as nn
from torch_concepts import Annotations, AnnotatedTensor
from torch_concepts.nn.modules.low.base.layer import BaseConceptLayer


class NormedConceptToConcept(BaseConceptLayer):
    """Layer-normalised linear concept predictor.

    Applies layer-norm to the input concept scores before the linear map.
    Useful when concept activations vary widely in scale.

    Args:
        in_concepts:  Input concept dimension (int or Annotations).
        out_concepts: Output concept dimension (int or Annotations).
    """

    def __init__(
        self,
        in_concepts:  Union[int, Annotations],
        out_concepts: Union[int, Annotations],
    ):
        super().__init__(in_concepts=in_concepts, out_concepts=out_concepts)
        self.norm      = nn.LayerNorm(self.in_concepts_shape)
        self.predictor = nn.Linear(self.in_concepts_shape, self.out_concepts_shape)

    def forward(self, concepts: torch.Tensor) -> torch.Tensor:
        x = self.norm(concepts)
        x = self.predictor(x)
        # If out_concepts is Annotations, wrap output as AnnotatedTensor
        return self.annotate(x)

Instantiation and forward pass:

import torch
import torch_concepts as pyc

annotations = pyc.Annotations(
    labels=["smoking", "genotype", "tar"],
    cardinalities=[1, 3, 1],
    types=["binary", "categorical", "continuous"],
)

# Plain int — no annotation on output
layer = NormedConceptToConcept(in_concepts=5, out_concepts=2)
out = layer(concepts=torch.randn(8, 5))   # (8, 2)

# Annotations — output is an AnnotatedTensor
layer = NormedConceptToConcept(in_concepts=annotations, out_concepts=2)
out = layer(concepts=torch.randn(8, 5))   # AnnotatedTensor (8, 2)
print(out["genotype"])                    # slice by name
Semantics-Aware Layers

Passing an Annotations object for in_concepts or out_concepts unlocks semantics-awareness:

  • self.in_concepts (or out_concepts) holds the full Annotations object.

  • self.in_concepts_shape still holds the corresponding integer (annotations.size).

  • Call self.annotate(x) (no arguments) to wrap the output as an AnnotatedTensor — this uses self.out_concepts if it is an Annotations.

Use the annotations inside the layer when you need per-concept branching, e.g. grouping columns by cardinality or type:

class TypeAwareConceptLayer(BaseConceptLayer):
    def __init__(self, in_concepts: Annotations, out_concepts: int):
        super().__init__(in_concepts=in_concepts, out_concepts=out_concepts)
        # Build one head per concept type
        binary_size = sum(
            c for c, t in zip(in_concepts.cardinalities, in_concepts.types)
            if t == "binary"
        )
        self.binary_head = nn.Linear(binary_size, out_concepts)
        # ... other heads

    def forward(self, concepts: torch.Tensor) -> torch.Tensor:
        ann: Annotations = self.in_concepts
        binary = concepts[:, [i for i, t in enumerate(ann.types) if t == "binary"]]
        return self.binary_head(binary)

MixConceptEmbeddingToConcept is the canonical example — it requires Annotations for in_concepts so it can split columns by cardinality and concept type.

Using a Layer inside a ParametricCPD

Layers integrate directly into the Mid-Level API through ParametricCPD. Pass your layer as the value under the distribution-parameter key that the target variable’s distribution expects:

import torch_concepts as pyc
from torch_concepts import ConceptVariable
from torch.distributions import Bernoulli
from torch_concepts.nn import ParametricCPD
from normed import NormedConceptToConcept   # the layer above

latent_var = pyc.EmbeddingVariable("latent", distribution=pyc.distributions.Delta, size=64)
concept    = ConceptVariable("smoking", distribution=Bernoulli)

cpd = ParametricCPD(
    concept,
    parents=[latent_var],
    parametrization={
        "logits": NormedConceptToConcept(in_concepts=64, out_concepts=1),
    },
)

The parametrization dict key must match the distribution’s constructor parameter (logits or probs for Bernoulli/OneHotCategorical, loc/scale for Normal).

Registering the Layer

  1. Add the file. Place your layer under the appropriate subdirectory:

    • torch_concepts/nn/modules/low/encoders/ — if it maps from embeddings

    • torch_concepts/nn/modules/low/predictors/ — if it maps from concept scores

    • Create a new subdirectory if neither fits.

  2. Export it. Add the class to torch_concepts/nn/__init__.py:

    from .modules.low.predictors.normed import NormedConceptToConcept
__all__ = [..., "NormedConceptToConcept"]

  3. Document it. Add an autoclass entry in doc/modules/low_level_api.rst:

    .. autoclass:: torch_concepts.nn.NormedConceptToConcept
   :members:

  4. Test it. Add a test in tests/nn/modules/low/ that checks the output shape for both int and Annotations inputs, and that AnnotatedTensor is returned when out_concepts is an Annotations.

Next Steps

  • See how layers compose into probabilistic models: Contributing a New Model.

  • Browse the Low-Level API reference: Low-Level API.

  • Check existing layer implementations in torch_concepts/nn/modules/low/.