temp fix

Dockerfile

@@ -3,7 +3,7 @@
 FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
 
 WORKDIR /app
-
+RUN ls .
 COPY . /app
 
 SHELL ["/bin/bash", "-c"]
@@ -11,12 +11,10 @@ SHELL ["/bin/bash", "-c"]
 USER root
 
 RUN ls /app
-RUN echo "---"
 RUN ls /app/src
-RUN echo "----"
-RUN ls /app/src/models/
-RUN echo "----"
-RUN ls /app/src/models/lgatr
+RUN ls /app/src/1models/
+RUN ls /app/src/1models/LGATr
+
 RUN apt update && \
     DEBIAN_FRONTEND=noninteractive apt install --yes --no-install-recommends \
         build-essential \

src/1models/GATr/Gatr.py

@@ -0,0 +1,104 @@
+from gatr import GATr, SelfAttentionConfig, MLPConfig
+from gatr.interface import (
+    embed_point,
+    extract_scalar,
+    extract_point,
+    embed_scalar,
+    embed_translation,
+    extract_translation
+)
+import torch
+import torch.nn as nn
+from xformers.ops.fmha import BlockDiagonalMask
+
+
+class GATrModel(torch.nn.Module):
+    def __init__(self, n_scalars, hidden_mv_channels, hidden_s_channels, blocks, embed_as_vectors, n_scalars_out):
+        super().__init__()
+        self.n_scalars = n_scalars
+        self.hidden_mv_channels = hidden_mv_channels
+        self.hidden_s_channels = hidden_s_channels
+        self.blocks = blocks
+        self.embed_as_vectors = embed_as_vectors
+        self.input_dim = 3
+        self.n_scalars_out = n_scalars_out
+        self.gatr = GATr(
+            in_mv_channels=1,
+            out_mv_channels=1,
+            hidden_mv_channels=hidden_mv_channels,
+            in_s_channels=n_scalars,
+            out_s_channels=n_scalars_out,
+            hidden_s_channels=hidden_s_channels,
+            num_blocks=blocks,
+            attention=SelfAttentionConfig(),  # Use default parameters for attention
+            mlp=MLPConfig(),  # Use default parameters for MLP
+        )
+        self.batch_norm = nn.BatchNorm1d(self.input_dim, momentum=0.1)
+        #self.clustering = nn.Linear(3, self.output_dim - 1, bias=False)
+        if n_scalars_out > 0:
+            self.beta = nn.Linear(n_scalars_out + 1, 1)
+        else:
+            self.beta = None
+
+    def forward(self, data):
+        # data: instance of EventBatch
+        inputs_v = data.input_vectors.float()
+        inputs_scalar = data.input_scalars.float()
+        assert inputs_scalar.shape[1] == self.n_scalars
+        if self.embed_as_vectors:
+            velocities = embed_translation(inputs_v)
+            embedded_inputs = (
+                velocities
+            )
+            # if it contains nans, raise an error
+            if torch.isnan(embedded_inputs).any():
+                raise ValueError("NaNs in the input!")
+        else:
+            inputs = inputs_v
+            embedded_inputs = embed_point(inputs)
+        embedded_inputs = embedded_inputs.unsqueeze(-2) # (batch_size*num_points, 1, 16)
+        mask = self.build_attention_mask(data.batch_idx)
+        embedded_outputs, output_scalars = self.gatr(
+            embedded_inputs, scalars=inputs_scalar, attention_mask=mask
+        )
+        #if self.embed_as_vectors:
+        #    x_clusters = extract_translation(embedded_outputs)
+        #else:
+        #    x_clusters = extract_point(embedded_outputs)
+        if self.embed_as_vectors:
+            x_clusters = extract_translation(embedded_outputs)
+        else:
+            x_clusters = extract_point(embedded_outputs)
+        original_scalar = extract_scalar(embedded_outputs)
+        if self.beta is not None:
+            beta = self.beta(torch.cat([original_scalar[:, 0, :], output_scalars], dim=1))
+            x = torch.cat((x_clusters[:, 0, :], torch.sigmoid(beta.view(-1, 1))), dim=1)
+        else:
+            x = x_clusters[:, 0, :]
+        if torch.isnan(x).any():
+            raise ValueError("NaNs in the output!")
+        #print(x[:5])
+        return x
+
+    def build_attention_mask(self, batch_numbers):
+        return BlockDiagonalMask.from_seqlens(
+            torch.bincount(batch_numbers.long()).tolist()
+        )
+
+def get_model(args, obj_score=False):
+    n_scalars_out = 8
+    if args.beta_type == "pt":
+        n_scalars_out = 0
+    elif args.beta_type == "pt+bc":
+        n_scalars_out = 8
+    n_scalars_in = 12
+    if args.no_pid:
+        n_scalars_in = 12-9
+    return GATrModel(
+        n_scalars=n_scalars_in,
+        hidden_mv_channels=args.hidden_mv_channels,
+        hidden_s_channels=args.hidden_s_channels,
+        blocks=args.num_blocks,
+        embed_as_vectors=args.embed_as_vectors,
+        n_scalars_out=n_scalars_out
+    )

src/1models/LGATr/lgatr.py

@@ -0,0 +1,196 @@
+from lgatr import GATr, SelfAttentionConfig, MLPConfig
+from lgatr.interface import embed_vector, extract_scalar, embed_spurions, extract_vector
+import torch
+import torch.nn as nn
+from xformers.ops.fmha import BlockDiagonalMask
+from torch_scatter import scatter_sum, scatter_max, scatter_mean
+
+
+class LGATrModel(torch.nn.Module):
+    def __init__(self, n_scalars, hidden_mv_channels, hidden_s_channels, blocks, embed_as_vectors, n_scalars_out, return_scalar_coords, obj_score=False, global_featuers_copy=False):
+        super().__init__()
+        self.return_scalar_coords = return_scalar_coords
+        self.n_scalars = n_scalars
+        self.hidden_mv_channels = hidden_mv_channels
+        self.hidden_s_channels = hidden_s_channels
+        self.blocks = blocks
+        self.embed_as_vectors = embed_as_vectors
+        self.input_dim = 3
+        self.n_scalars_out = n_scalars_out
+        self.obj_score = obj_score
+        self.global_features_copy = global_featuers_copy
+        self.gatr = GATr(
+            in_mv_channels=3,
+            out_mv_channels=1,
+            hidden_mv_channels=hidden_mv_channels,
+            in_s_channels=n_scalars,
+            out_s_channels=n_scalars_out,
+            hidden_s_channels=hidden_s_channels,
+            num_blocks=blocks,
+            attention=SelfAttentionConfig(),  # Use default parameters for attention
+            mlp=MLPConfig(),  # Use default parameters for MLP
+        )
+        if self.global_features_copy:
+            self.gatr_global_features = GATr(
+                in_mv_channels=3,
+                out_mv_channels=1,
+                hidden_mv_channels=hidden_mv_channels,
+                in_s_channels=n_scalars,
+                out_s_channels=n_scalars_out,
+                hidden_s_channels=hidden_s_channels,
+                num_blocks=blocks,
+                attention=SelfAttentionConfig(),  # Use default parameters for attention
+                mlp=MLPConfig(),  # Use default parameters for MLP
+            )
+        #self.batch_norm = nn.BatchNorm1d(self.input_dim, momentum=0.1)
+        #self.clustering = nn.Linear(3, self.output_dim - 1, bias=False)
+        if n_scalars_out > 0:
+            if obj_score:
+                factor = 1
+                if self.global_features_copy: factor = 2
+                self.beta = nn.Sequential(
+                    nn.Linear((n_scalars_out + 1) * factor, 10),
+                    nn.LeakyReLU(),
+                    nn.Linear(10, 1),
+                    #nn.Sigmoid()
+                )
+            else:
+                self.beta = nn.Linear(n_scalars_out + 1, 1)
+        else:
+            self.beta = None
+
+    def forward(self, data, data_events=None, data_events_clusters=None, cpu_demo=False):
+        # data: instance of EventBatch
+        if self.global_features_copy:
+            assert data_events is not None and data_events_clusters is not None
+            assert self.obj_score
+            inputs_v = data_events.input_vectors
+            inputs_scalar = data_events.input_scalars
+            assert inputs_scalar.shape[1] == self.n_scalars, "Expected %d, got %d" % (
+            self.n_scalars, inputs_scalar.shape[1])
+            mask_global = self.build_attention_mask(data_events.batch_idx)
+            embedded_inputs_events = embed_vector(inputs_v.unsqueeze(0))
+            multivectors = embedded_inputs_events.unsqueeze(-2)
+            spurions = embed_spurions(beam_reference="xyplane", add_time_reference=True,
+                                      device=multivectors.device, dtype=multivectors.dtype)
+
+            num_points, x = inputs_v.shape
+            assert x == 4
+            spurions = spurions[None, None, ...].repeat(1, num_points, 1, 1)  # (batchsize, num_points, 2, 16)
+            multivectors = torch.cat((multivectors, spurions), dim=-2)
+            embedded_outputs, output_scalars = self.gatr_global_features(
+                multivectors, scalars=inputs_scalar, attention_mask=mask_global
+            )
+            original_scalar = extract_scalar(embedded_outputs)
+            scalar_embeddings_nodes = torch.cat([original_scalar[0, :, 0, :], output_scalars[0, :, :]], dim=1)
+            scalar_embeddings_global = scatter_mean(scalar_embeddings_nodes, torch.tensor(data_events_clusters).to(scalar_embeddings_nodes.device)+1, dim=0)[1:]
+
+        inputs_v = data.input_vectors.float() # four-momenta
+        inputs_scalar = data.input_scalars.float()
+        assert inputs_scalar.shape[1] == self.n_scalars
+        num_points, x = inputs_v.shape
+        assert x == 4
+        #velocities = embed_vector(inputs_v)
+
+        inputs_v = inputs_v.unsqueeze(0)
+        embedded_inputs = embed_vector(inputs_v)
+        # if it contains nans, raise an error
+        if torch.isnan(embedded_inputs).any():
+            raise ValueError("NaNs in the input!")
+        multivectors = embedded_inputs.unsqueeze(-2) # (batch_size*num_points, 1, 16)
+        # for spurions, duplicate each unique batch_idx. e.g. [0,0,1,1,2,2] etc.
+        #spurions_batch_idx = torch.repeat_interleave(data.batch_idx.unique(), 2)
+        #batch_idx = torch.cat([data.batch_idx, spurions_batch_idx])
+        spurions = embed_spurions(beam_reference="xyplane", add_time_reference=True,
+                                  device=multivectors.device, dtype=multivectors.dtype)
+        spurions = spurions[None, None, ...].repeat(1, num_points, 1, 1)  # (batchsize, num_points, 2, 16)
+        multivectors = torch.cat((multivectors, spurions), dim=-2)  # (batchsize, num_points, 3, 16) - Just embed the spurions as two extra multivector channels
+        mask = self.build_attention_mask(data.batch_idx)
+        if cpu_demo:
+            mask = None
+        embedded_outputs, output_scalars = self.gatr(
+            multivectors, scalars=inputs_scalar, attention_mask=mask
+        )
+
+        #if self.embed_as_vectors:
+        #    x_clusters = extract_translation(embedded_outputs)
+        #else:
+        #    x_clusters = extract_point(embedded_outputs)
+        x_clusters = extract_vector(embedded_outputs)
+        original_scalar = extract_scalar(embedded_outputs)
+        if self.beta is not None:
+            if self.obj_score:
+                extract_from_virtual_nodes = False
+                # assert that data has fake_nodes_idx from which we read the objectness score
+                #assert "fake_nodes_idx" in data.__dict__
+                # print batch number 3 and 4 inputs
+                #for nbatch in [3, 4]:
+                #    print("#### batch no. ", nbatch , "#######")
+               #     print(" -> scalar inputs", inputs_scalar[data.batch_idx==nbatch].shape, inputs_scalar[data.batch_idx == nbatch])
+               #     print(" -> vector inputs", data.input_vectors[data.batch_idx==nbatch].shape, data.input_vectors[data.batch_idx == nbatch])
+               #     print("############")
+                scalar_embeddings = torch.cat([original_scalar[0, :, 0, :], output_scalars[0, :, :]], dim=1)
+                if extract_from_virtual_nodes:
+                    values = torch.cat([original_scalar[0, data.fake_nodes_idx, 0, :], output_scalars[0, data.fake_nodes_idx, :]], dim=1)
+                else:
+                    values = scatter_mean(scalar_embeddings, data.batch_idx.to(scalar_embeddings.device).long(), dim=0)
+                if self.global_features_copy:
+                    values = torch.cat([values, scalar_embeddings_global], dim=1)
+                beta = self.beta(values)
+                #beta = self.beta(values)
+                return beta
+            vals = torch.cat([original_scalar[0, :, 0, :], output_scalars[0, :, :]], dim=1)
+            beta = self.beta(vals)
+            if self.return_scalar_coords:
+                x = output_scalars[0, :, :3]
+                #print(x.shape)
+                #print(x[:5])
+                x = torch.cat((x, torch.sigmoid(beta.view(-1, 1))), dim=1)
+            else:
+                x = torch.cat((x_clusters[0, :, 0, :], torch.sigmoid(beta.view(-1, 1))), dim=1)
+        else:
+            x = x_clusters[:, 0, :]
+        if torch.isnan(x).any():
+            raise ValueError("NaNs in the output!")
+        #print(x[:5])
+        print("LGATr x shape:", x.shape)
+        return x
+
+    def build_attention_mask(self, batch_numbers):
+        return BlockDiagonalMask.from_seqlens(
+            torch.bincount(batch_numbers.long()).tolist()
+        )
+
+def get_model(args, obj_score=False):
+    n_scalars_out = 8
+    if args.beta_type == "pt":
+        n_scalars_out = 0
+    elif args.beta_type == "pt+bc":
+        n_scalars_out = 8
+    n_scalars_in = 12
+    if args.no_pid:
+        n_scalars_in = 12 - 9
+    if obj_score:
+        return LGATrModel(
+            n_scalars=n_scalars_in,
+            hidden_mv_channels=8,
+            hidden_s_channels=16,
+            blocks=5,
+            embed_as_vectors=False,
+            n_scalars_out=n_scalars_out,
+            return_scalar_coords=args.scalars_oc,
+            obj_score=obj_score,
+            global_featuers_copy=args.global_features_obj_score
+        )
+
+    return LGATrModel(
+        n_scalars=n_scalars_in,
+        hidden_mv_channels=args.hidden_mv_channels,
+        hidden_s_channels=args.hidden_s_channels,
+        blocks=args.num_blocks,
+        embed_as_vectors=args.embed_as_vectors,
+        n_scalars_out=n_scalars_out,
+        return_scalar_coords=args.scalars_oc,
+        obj_score=obj_score
+    )
+

src/1models/identity.py

@@ -0,0 +1,22 @@
+import torch
+
+class IdentityModel(torch.nn.Module):
+    def __init__(self, n_out_coords=3):
+        super().__init__()
+        self.n_out_coords = n_out_coords
+
+    def forward(self, data):
+        # data: instance of EventBatch
+        inputs_v = data.input_vectors # four-momenta
+        betas = torch.ones(data.input_vectors.shape[0]).to(inputs_v.device)
+        norm_inputs_v = torch.norm(inputs_v, dim=1).unsqueeze(1)
+        #print("inputs_v.shape", inputs_v.shape)
+        #print("betas.shape", betas.shape)
+        #print("norm_inputs_v.shape", norm_inputs_v.shape)
+        #print("betas unsqueezed shape", betas.unsqueeze(1).shape)
+        x = torch.cat([inputs_v / norm_inputs_v, betas.unsqueeze(1)], dim=1)
+        return x
+
+
+def get_model(args):
+    return IdentityModel()

src/1models/transformer/tr_blocks.py

@@ -0,0 +1,531 @@
+# File copied from https://raw.githubusercontent.com/heidelberg-hepml/lorentz-gatr/refs/heads/main/experiments/baselines/transformer.py
+from functools import partial
+from typing import Optional, Tuple
+import torch
+from einops import rearrange
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from lgatr.layers import ApplyRotaryPositionalEncoding
+from lgatr.primitives.attention import scaled_dot_product_attention
+
+
+def to_nd(tensor, d):
+    """Make tensor n-dimensional, group extra dimensions in first."""
+    return tensor.view(
+        -1, *(1,) * (max(0, d - 1 - tensor.dim())), *tensor.shape[-(d - 1) :]
+    )
+
+class BaselineLayerNorm(nn.Module):
+    """Baseline layer norm over all dimensions except the first."""
+
+    @staticmethod
+    def forward(inputs: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor
+            Input data
+
+        Returns
+        -------
+        outputs : Tensor
+            Normalized inputs.
+        """
+        return torch.nn.functional.layer_norm(
+            inputs, normalized_shape=inputs.shape[-1:]
+        )
+
+
+class MultiHeadQKVLinear(nn.Module):
+    """Compute queries, keys, and values via multi-head attention.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    hidden_channels : int
+        Number of hidden channels = size of query, key, and value.
+    num_heads : int
+        Number of attention heads.
+    """
+
+    def __init__(self, in_channels, hidden_channels, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.linear = nn.Linear(in_channels, 3 * hidden_channels * num_heads)
+
+    def forward(self, inputs):
+        """Forward pass.
+
+        Returns
+        -------
+        q : Tensor
+            Queries
+        k : Tensor
+            Keys
+        v : Tensor
+            Values
+        """
+        qkv = self.linear(inputs)  # (..., num_items, 3 * hidden_channels * num_heads)
+        q, k, v = rearrange(
+            qkv,
+            "... items (qkv hidden_channels num_heads) -> qkv ... num_heads items hidden_channels",
+            num_heads=self.num_heads,
+            qkv=3,
+        )
+        return q, k, v
+
+
+class MultiQueryQKVLinear(nn.Module):
+    """Compute queries, keys, and values via multi-query attention.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    hidden_channels : int
+        Number of hidden channels = size of query, key, and value.
+    num_heads : int
+        Number of attention heads.
+    """
+
+    def __init__(self, in_channels, hidden_channels, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.q_linear = nn.Linear(in_channels, hidden_channels * num_heads)
+        self.k_linear = nn.Linear(in_channels, hidden_channels)
+        self.v_linear = nn.Linear(in_channels, hidden_channels)
+
+    def forward(self, inputs):
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor
+            Input data
+
+        Returns
+        -------
+        q : Tensor
+            Queries
+        k : Tensor
+            Keys
+        v : Tensor
+            Values
+        """
+        q = rearrange(
+            self.q_linear(inputs),
+            "... items (hidden_channels num_heads) -> ... num_heads items hidden_channels",
+            num_heads=self.num_heads,
+        )
+        k = self.k_linear(inputs)[
+            ..., None, :, :
+        ]  # (..., head=1, item, hidden_channels)
+        v = self.v_linear(inputs)[..., None, :, :]
+        return q, k, v
+
+
+class BaselineSelfAttention(nn.Module):
+    """Baseline self-attention layer.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of input channels.
+    hidden_channels : int
+        Number of hidden channels = size of query, key, and value.
+    num_heads : int
+        Number of attention heads.
+    pos_encoding : bool
+        Whether to apply rotary positional embeddings along the item dimension to the scalar keys
+        and queries.
+    pos_enc_base : int
+        Maximum frequency used in positional encodings. (The minimum frequency is always 1.)
+    multi_query : bool
+        Use multi-query attention instead of multi-head attention.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        num_heads: int = 8,
+        pos_encoding: bool = False,
+        pos_enc_base: int = 4096,
+        multi_query: bool = True,
+        dropout_prob=None,
+    ) -> None:
+        super().__init__()
+
+        # Store settings
+        self.num_heads = num_heads
+        self.hidden_channels = hidden_channels
+
+        # Linear maps
+        qkv_class = MultiQueryQKVLinear if multi_query else MultiHeadQKVLinear
+        self.qkv_linear = qkv_class(in_channels, hidden_channels, num_heads)
+        self.out_linear = nn.Linear(hidden_channels * num_heads, out_channels)
+
+        # Optional positional encoding
+        if pos_encoding:
+            self.pos_encoding = ApplyRotaryPositionalEncoding(
+                hidden_channels, item_dim=-2, base=pos_enc_base
+            )
+        else:
+            self.pos_encoding = None
+
+        if dropout_prob is not None:
+            self.dropout = nn.Dropout(dropout_prob)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor
+            Input data
+        attention_mask : None or Tensor or xformers.ops.AttentionBias
+            Optional attention mask
+
+        Returns
+        -------
+        outputs : Tensor
+            Outputs
+        """
+        q, k, v = self.qkv_linear(
+            inputs
+        )  # each: (..., num_heads, num_items, num_channels, 16)
+        # Rotary positional encoding
+        if self.pos_encoding is not None:
+            q = self.pos_encoding(q)
+            k = self.pos_encoding(k)
+
+        # Attention layer
+        h = self._attend(q, k, v, attention_mask, is_causal=is_causal)
+
+        # Concatenate heads and transform linearly
+        h = rearrange(
+            h,
+            "... num_heads num_items hidden_channels -> ... num_items (num_heads hidden_channels)",
+        )
+        outputs = self.out_linear(h)  # (..., num_items, out_channels)
+
+        if self.dropout is not None:
+            outputs = self.dropout(outputs)
+
+        return outputs
+
+    @staticmethod
+    def _attend(q, k, v, attention_mask=None, is_causal=False):
+        """Scaled dot-product attention."""
+
+        # Add batch dimension if needed
+        bh_shape = q.shape[:-2]
+        q = to_nd(q, 4)
+        k = to_nd(k, 4)
+        v = to_nd(v, 4)
+
+        # SDPA
+        outputs = scaled_dot_product_attention(
+            q.contiguous(),
+            k.expand_as(q).contiguous(),
+            v.expand_as(q).contiguous(),
+            attn_mask=attention_mask,
+            is_causal=is_causal,
+        )
+
+        # Return batch dimensions to inputs
+        outputs = outputs.view(*bh_shape, *outputs.shape[-2:])
+
+        return outputs
+
+
+class BaselineTransformerBlock(nn.Module):
+    """Baseline transformer block.
+
+    Inputs are first processed by a block consisting of LayerNorm, multi-head self-attention, and
+    residual connection. Then the data is processed by a block consisting of another LayerNorm, an
+    item-wise two-layer MLP with GeLU activations, and another residual connection.
+
+    Parameters
+    ----------
+    channels : int
+        Number of input and output channels.
+    num_heads : int
+        Number of attention heads.
+    pos_encoding : bool
+        Whether to apply rotary positional embeddings along the item dimension to the scalar keys
+        and queries.
+    pos_encoding_base : int
+        Maximum frequency used in positional encodings. (The minimum frequency is always 1.)
+    increase_hidden_channels : int
+        Factor by which the key, query, and value size is increased over the default value of
+        hidden_channels / num_heads.
+    multi_query : bool
+        Use multi-query attention instead of multi-head attention.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads: int = 8,
+        pos_encoding: bool = False,
+        pos_encoding_base: int = 4096,
+        increase_hidden_channels=1,
+        multi_query: bool = True,
+        dropout_prob=None,
+    ) -> None:
+        super().__init__()
+
+        self.norm = BaselineLayerNorm()
+
+        # When using positional encoding, the number of scalar hidden channels needs to be even.
+        # It also should not be too small.
+        hidden_channels = channels // num_heads * increase_hidden_channels
+        if pos_encoding:
+            hidden_channels = (hidden_channels + 1) // 2 * 2
+            hidden_channels = max(hidden_channels, 16)
+
+        self.attention = BaselineSelfAttention(
+            channels,
+            channels,
+            hidden_channels,
+            num_heads=num_heads,
+            pos_encoding=pos_encoding,
+            pos_enc_base=pos_encoding_base,
+            multi_query=multi_query,
+            dropout_prob=dropout_prob,
+        )
+
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, 2 * channels),
+            nn.Dropout(dropout_prob) if dropout_prob is not None else nn.Identity(),
+            nn.GELU(),
+            nn.Linear(2 * channels, channels),
+            nn.Dropout(dropout_prob) if dropout_prob is not None else nn.Identity(),
+        )
+
+    def forward(
+        self, inputs: torch.Tensor, attention_mask=None, is_causal=False
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor
+            Input data
+        attention_mask : None or Tensor or xformers.ops.AttentionBias
+            Optional attention mask
+
+        Returns
+        -------
+        outputs : Tensor
+            Outputs
+        """
+
+        # Residual attention
+        h = self.norm(inputs)
+        h = self.attention(h, attention_mask=attention_mask, is_causal=is_causal)
+        outputs = inputs + h
+
+        # Residual MLP
+        h = self.norm(outputs)
+        h = self.mlp(h)
+        outputs = outputs + h
+
+        return outputs
+
+
+class Transformer(nn.Module):
+    """Baseline transformer.
+
+    Combines num_blocks transformer blocks, each consisting of multi-head self-attention layers, an
+    MLP, residual connections, and normalization layers.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    hidden_channels : int
+        Number of hidden channels.
+    num_blocks : int
+        Number of transformer blocks.
+    num_heads : int
+        Number of attention heads.
+    pos_encoding : bool
+        Whether to apply rotary positional embeddings along the item dimension to the scalar keys
+        and queries.
+    pos_encoding_base : int
+        Maximum frequency used in positional encodings. (The minimum frequency is always 1.)
+    increase_hidden_channels : int
+        Factor by which the key, query, and value size is increased over the default value of
+        hidden_channels / num_heads.
+    multi_query : bool
+        Use multi-query attention instead of multi-head attention.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        num_blocks: int = 10,
+        num_heads: int = 8,
+        pos_encoding: bool = False,
+        pos_encoding_base: int = 4096,
+        checkpoint_blocks: bool = False,
+        increase_hidden_channels=1,
+        multi_query: bool = False,
+        dropout_prob=None,
+    ) -> None:
+        super().__init__()
+        self.checkpoint_blocks = checkpoint_blocks
+        self.linear_in = nn.Linear(in_channels, hidden_channels)
+        self.blocks = nn.ModuleList(
+            [
+                BaselineTransformerBlock(
+                    hidden_channels,
+                    num_heads=num_heads,
+                    pos_encoding=pos_encoding,
+                    pos_encoding_base=pos_encoding_base,
+                    increase_hidden_channels=increase_hidden_channels,
+                    multi_query=multi_query,
+                    dropout_prob=dropout_prob,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+        self.linear_out = nn.Linear(hidden_channels, out_channels)
+
+    def forward(
+        self, inputs: torch.Tensor, attention_mask=None, is_causal=False
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor with shape (..., num_items, num_channels)
+            Input data
+        attention_mask : None or Tensor or xformers.ops.AttentionBias
+            Optional attention mask
+        is_causal: bool
+
+        Returns
+        -------
+        outputs : Tensor with shape (..., num_items, num_channels)
+            Outputs
+        """
+        h = self.linear_in(inputs)
+        for block in self.blocks:
+            if self.checkpoint_blocks:
+                fn = partial(block, attention_mask=attention_mask, is_causal=is_causal)
+                h = checkpoint(fn, h)
+            else:
+                h = block(h, attention_mask=attention_mask, is_causal=is_causal)
+        outputs = self.linear_out(h)
+        return outputs
+
+
+class AxialTransformer(nn.Module):
+    """Baseline axial transformer for data with two token dimensions.
+
+    Combines num_blocks transformer blocks, each consisting of multi-head self-attention layers, an
+    MLP, residual connections, and normalization layers.
+
+    Assumes input data with shape `(..., num_items_1, num_items_2, num_channels, [16])`.
+
+    The first, third, fifth, ... block computes attention over the `items_2` axis. The other blocks
+    compute attention over the `items_1` axis. Positional encoding can be specified separately for
+    both axes.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    hidden_channels : int
+        Number of hidden channels.
+    num_blocks : int
+        Number of transformer blocks.
+    num_heads : int
+        Number of attention heads.
+    pos_encodings : tuple of bool
+        Whether to apply rotary positional embeddings along the item dimensions to the scalar keys
+        and queries.
+    pos_encoding_base : int
+        Maximum frequency used in positional encodings. (The minimum frequency is always 1.)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        num_blocks: int = 20,
+        num_heads: int = 8,
+        pos_encodings: Tuple[bool, bool] = (False, False),
+        pos_encoding_base: int = 4096,
+    ) -> None:
+        super().__init__()
+        self.linear_in = nn.Linear(in_channels, hidden_channels)
+        self.blocks = nn.ModuleList(
+            [
+                BaselineTransformerBlock(
+                    hidden_channels,
+                    num_heads=num_heads,
+                    pos_encoding=pos_encodings[(block + 1) % 2],
+                    pos_encoding_base=pos_encoding_base,
+                )
+                for block in range(num_blocks)
+            ]
+        )
+        self.linear_out = nn.Linear(hidden_channels, out_channels)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        inputs : Tensor with shape (..., num_items1, num_items2, num_channels)
+            Input data
+
+        Returns
+        -------
+        outputs : Tensor with shape (..., num_items1, num_items2, num_channels)
+            Outputs
+        """
+
+        rearrange_pattern = "... i j c -> ... j i c"
+
+        h = self.linear_in(inputs)
+
+        for i, block in enumerate(self.blocks):
+            # For first, third, ... block, we want to perform attention over the first token
+            # dimension. We implement this by transposing the two item dimensions.
+            if i % 2 == 1:
+                h = rearrange(h, rearrange_pattern)
+
+            h = block(h)
+
+            # Transposing back to standard axis order
+            if i % 2 == 1:
+                h = rearrange(h, rearrange_pattern)
+
+        outputs = self.linear_out(h)
+
+        return outputs

src/1models/transformer/transformer.py

@@ -0,0 +1,141 @@
+from src.models.transformer.tr_blocks import Transformer
+import torch
+import torch.nn as nn
+from xformers.ops.fmha import BlockDiagonalMask
+from torch_scatter import scatter_max, scatter_add, scatter_mean
+import numpy as np
+
+
+class TransformerModel(torch.nn.Module):
+    def __init__(self, n_scalars, n_scalars_out, n_blocks, n_heads, internal_dim, obj_score, global_features_copy=False):
+        super().__init__()
+        self.n_scalars = n_scalars
+        self.input_dim = n_scalars + 3
+        if obj_score:
+            self.input_dim += 1
+        self.output_dim = 3
+        self.obj_score = obj_score
+        #internal_dim = 128
+        #self.custom_decoder = nn.Linear(internal_dim, self.output_dim)
+        #n_heads = 4
+        #self.transformer = nn.TransformerEncoder(
+        #    nn.TransformerEncoderLayer(
+        #        d_model=n_heads*self.input_dim,
+        #        nhead=n_heads,
+        #        dim_feedforward=internal_dim,
+        #        dropout=0.1,
+        #        activation="gelu",
+        #    ),
+        #    num_layers=4,
+        #)
+        if n_scalars_out > 0:
+            self.output_dim += 1 # betas regression
+        if self.obj_score:
+            self.output_dim = 10
+        self.global_features_copy = global_features_copy
+        self.transformer = Transformer(
+            in_channels=self.input_dim,
+            out_channels=self.output_dim,
+            hidden_channels=internal_dim,
+            num_heads=n_heads,
+            num_blocks=n_blocks,
+        )
+        if self.global_features_copy:
+            self.transformer_global_features = Transformer(
+                in_channels=self.input_dim,
+                out_channels=self.output_dim,
+                hidden_channels=internal_dim,
+                num_heads=n_heads,
+                num_blocks=n_blocks,
+            )
+        self.batch_norm = nn.BatchNorm1d(self.input_dim, momentum=0.1)
+        if self.obj_score:
+            factor = 1
+            if self.global_features_copy: factor = 2
+            self.final_mlp = nn.Sequential(
+                nn.Linear(self.output_dim*factor, 10),
+                nn.LeakyReLU(),
+                nn.Linear(10, 1),
+            )
+        #self.clustering = nn.Linear(3, self.output_dim - 1, bias=False)
+
+    def forward(self, data, data_events=None, data_events_clusters=None):
+        # data: instance of EventBatch
+        # data_events & data_events_clusters: Only relevant if --global-features-obj-score is on: data_events contains
+        # the "unmodified" batch where the batch indices are
+        if self.global_features_copy:
+            assert data_events is not None and data_events_clusters is not None
+            assert self.obj_score
+            inputs_v = data_events.input_vectors.float()
+            inputs_scalar = data_events.input_scalars.float()
+            assert inputs_scalar.shape[1] == self.n_scalars, "Expected %d, got %d" % (
+            self.n_scalars, inputs_scalar.shape[1])
+            inputs_transformer_events = torch.cat([inputs_scalar, inputs_v], dim=1)
+            inputs_transformer_events = inputs_transformer_events.float()
+            assert inputs_transformer_events.shape[1] == self.input_dim
+            mask_global = self.build_attention_mask(data_events.batch_idx)
+            x_global = inputs_transformer_events.unsqueeze(0)
+            x_global = self.transformer_global_features(x_global, attention_mask=mask_global)[0]
+            assert x_global.shape[1] == self.output_dim, "Expected %d, got %d" % (self.output_dim, x_global.shape[1])
+            assert x_global.shape[0] == x_global.shape[0], "Expected %d, got %d" % (
+            inputs_transformer_events.shape[0], x_global.shape[0])
+            m_global = scatter_mean(x_global, torch.tensor(data_events_clusters).to(x_global.device)+1, dim=0)[1:]
+        inputs_v = data.input_vectors
+        inputs_scalar = data.input_scalars
+        assert inputs_scalar.shape[1] == self.n_scalars, "Expected %d, got %d" % (self.n_scalars, inputs_scalar.shape[1])
+        inputs_transformer = torch.cat([inputs_scalar, inputs_v], dim=1)
+        inputs_transformer = inputs_transformer.float()
+        print("input_dim", self.input_dim, inputs_transformer.shape)
+        assert inputs_transformer.shape[1] == self.input_dim
+        mask = self.build_attention_mask(data.batch_idx)
+        x = inputs_transformer.unsqueeze(0)
+        x = self.transformer(x, attention_mask=mask)[0]
+        assert x.shape[1] == self.output_dim, "Expected %d, got %d" % (self.output_dim, x.shape[1])
+        assert x.shape[0] == inputs_transformer.shape[0], "Expected %d, got %d" % (inputs_transformer.shape[0], x.shape[0])
+        if not self.obj_score:
+            x[:, -1] = torch.sigmoid(x[:, -1])
+        else:
+            extract_from_virtual_nodes = False
+            if extract_from_virtual_nodes:
+                x = self.final_mlp(x[data.fake_nodes_idx]) # x is the raw logits
+            else:
+                m = scatter_mean(x, torch.tensor(data.batch_idx).long().to(x.device), dim=0)
+                assert not "fake_nodes_idx" in data.__dict__
+                if self.global_features_copy:
+                    m = torch.cat([m, m_global], dim=1)
+                x = self.final_mlp(m).flatten()
+        return x
+
+    def build_attention_mask(self, batch_numbers):
+        return BlockDiagonalMask.from_seqlens(
+            torch.bincount(batch_numbers.long()).tolist()
+        )
+
+def get_model(args, obj_score=False):
+    n_scalars_out = 8
+    if args.beta_type == "pt":
+        n_scalars_out = 0
+    elif args.beta_type == "pt+bc":
+        n_scalars_out = 1
+    n_scalars_in = 12
+    if args.no_pid:
+        n_scalars_in = 12-9
+    if obj_score:
+        return TransformerModel(
+            n_scalars=n_scalars_in,
+            n_scalars_out=10,
+            n_blocks=5,
+            n_heads=args.n_heads,
+            internal_dim=64,
+            obj_score=obj_score,
+            global_features_copy=args.global_features_obj_score
+        )
+    return TransformerModel(
+        n_scalars=n_scalars_in,
+        n_scalars_out=n_scalars_out,
+        n_blocks=args.num_blocks,
+        n_heads=args.n_heads,
+        internal_dim=args.internal_dim,
+        obj_score=obj_score
+    )
+

src/model_wrapper_gradio.py

@@ -41,7 +41,7 @@ def inference(loss_str, train_dataset_str, input_text, input_text_quarks):
     args.spatial_part_only = True # LGATr
     args.load_model_weights =  model_path
     args.aug_soft = True # LGATr_GP etc.
-    args.network_config = "src/models/LGATr/lgatr.py"
+    args.network_config = "src/1models/LGATr/lgatr.py"
     args.beta_type = "pt+bc"
     args.embed_as_vectors = False
     args.debug = False