reasoning
The correctness test failed. There must be an issue with the attention kernel. Looking at the code, I see a potential problem - the attention kernel has a different shape handling. The indices for accessing k and v in the attention kernel are wrong.
Looking at the reshaping:
- q_reshaped = q_out.view({B, S, N, h, d}).permute({0, 1, 3, 2, 4}).contiguous()
This gives [B, S, h, N, d]
- k_reshaped and v_reshaped are similar.
But in the attention kernel, I'm computing:
- k_idx = b * S * h * N * d + s * h * N * d + head * N * d + key_pos * d + i
This seems correct for the [B, S, h, N, d] layout, where the stride is:
- For B: S * h * N * d
- For S: h * N * d
- For h: N * d
- For N: d
- For d: 1
So for element [b, s, head, key_pos, i], the index is: b * (S * h * N * d) + s * (h * N * d) + head * (N * d) + key_pos * d + i
Wait, but the attention is computing q @ k^T where k needs to be transposed. In the current implementation, for each position (b, s, head), we're iterating over all key positions to compute attention scores. But the way I'm indexing k is:
k_idx = b * S * h * N * d + s * h * N * d + head * N * d + key_pos * d + i
This accesses element [b, s, head, key_pos, i] in the k tensor. This is correct!
The issue might be in the softmax computation. Let me check - I'm only doing softmax on one thread (threadIdx.x == 0), which means only one thread updates the attn_smem. But all threads read from it. The __syncthreads() should ensure correctness.
Actually, I think the issue might be that I'm not handling the transpose correctly. In attention:
- q has shape [B, S, h, N, d] - query for each position
- k has shape [B, S, h, N, d] - key for each position
- We compute q @ k^T
For a single (b, s, head) combination:
- q for this position is [N, d]
- k for this position is [N, d]
- We compute [N, d] @ [d, N] = [N, N]
But wait, the reference code does:
q = self.q(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4) # [B, S, h, N, d]
k = self.k(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4) # [B, S, h, N, d]
attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)
So q has shape [B, S, h, N, d] and k has shape [B, S, h, N, d].
q @ k.transpose(-1, -2) gives [B, S, h, N, N].
For each (b, s, head), we compute:
q[b, s, head] @ k[b, s, head]^T = [N, d] @ [d, N] = [N, N]
This is for self-attention within each sequence position, across all residues N. But my kernel is computing attention for each (b, s, head) position, but I think there's a confusion.
Wait, I think I'm misunderstanding the structure. Let me re-read the reference:
- msa has shape [B, S, N, D] where B=batch, S=MSA depth, N=sequence length, D=dimension
- After layer norm, we compute q, k, v projections
- q = q_weight(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4) gives [B, S, h, N, d]
So for MSA row self-attention:
- We have B batches, S MSA sequences, h heads
- For each batch and MSA sequence, we attend over N positions
- The attention is: for position i in N, we compute attention weights over all positions j in N
So the attention computation should be:
- For each (b, s, head):
- q has shape [N, d]
- k has shape [N, d]
- scores = q @ k^T has shape [N, N]
- For each query position q_idx, we attend to all key positions
In my kernel, I'm handling one (b, s, head) per block. But I'm computing attention for the q at position q_idx for all keys. Wait, I don't see where I'm setting q_idx!
Looking at my code again:
```
int head_idx = b * S * h * N * d + s * h * N * d + head * N * d;
```
This is the base index for the entire [N, d] array for this (b, s, head). I then load q_smem with the entire [N, d] array? No, I'm only loading d elements into q_smem, which corresponds to position 0 in N.
This is wrong! I need to compute attention for ALL N query positions, not just position 0.
The reference computes full [N, N] attention matrix. Let me fix this.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom torch.utils.cpp_extension import load_inline\n\ncpp_source = r\"\"\"\n#include <torch/extension.h>\n#include <vector>\n\nstd::vector<torch::Tensor> msa_row_attn_cuda(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor q_weight,\n torch::Tensor k_weight,\n torch::Tensor v_weight,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n int num_heads,\n double eps);\n\ntorch::Tensor outer_product_mean_cuda(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor a_weight,\n torch::Tensor a_bias,\n torch::Tensor b_weight,\n torch::Tensor b_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n double eps);\n\ntorch::Tensor tri_mul_out_cuda(\n torch::Tensor pair,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor pa_weight,\n torch::Tensor pa_bias,\n torch::Tensor pb_weight,\n torch::Tensor pb_bias,\n torch::Tensor ga_weight,\n torch::Tensor ga_bias,\n torch::Tensor gb_weight,\n torch::Tensor gb_bias,\n torch::Tensor norm_out_weight,\n torch::Tensor norm_out_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n torch::Tensor gate_weight,\n torch::Tensor gate_bias,\n double eps);\n\nstd::vector<torch::Tensor> msa_row_attn(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor q_weight,\n torch::Tensor k_weight,\n torch::Tensor v_weight,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n int num_heads,\n double eps) {\n return msa_row_attn_cuda(msa, ln_weight, ln_bias, q_weight, k_weight, v_weight, out_weight, out_bias, num_heads, eps);\n}\n\ntorch::Tensor outer_product_mean(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor a_weight,\n torch::Tensor a_bias,\n torch::Tensor b_weight,\n torch::Tensor b_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n double eps) {\n return outer_product_mean_cuda(msa, ln_weight, ln_bias, a_weight, a_bias, b_weight, b_bias, out_weight, out_bias, eps);\n}\n\ntorch::Tensor tri_mul_out(\n torch::Tensor pair,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor pa_weight,\n torch::Tensor pa_bias,\n torch::Tensor pb_weight,\n torch::Tensor pb_bias,\n torch::Tensor ga_weight,\n torch::Tensor ga_bias,\n torch::Tensor gb_weight,\n torch::Tensor gb_bias,\n torch::Tensor norm_out_weight,\n torch::Tensor norm_out_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n torch::Tensor gate_weight,\n torch::Tensor gate_bias,\n double eps) {\n return tri_mul_out_cuda(pair, ln_weight, ln_bias, pa_weight, pa_bias, pb_weight, pb_bias, \n ga_weight, ga_bias, gb_weight, gb_bias, norm_out_weight, norm_out_bias,\n out_weight, out_bias, gate_weight, gate_bias, eps);\n}\n\"\"\"\n\ncuda_source = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\nconstexpr int BLOCK_SIZE = 256;\n\n__global__ void layernorm_multi_linear_kernel(\n const float* __restrict__ input,\n const float* __restrict__ ln_weight,\n const float* __restrict__ ln_bias,\n float* __restrict__ q_out,\n float* __restrict__ k_out,\n float* __restrict__ v_out,\n const float* __restrict__ q_weight,\n const float* __restrict__ k_weight,\n const float* __restrict__ v_weight,\n int num_groups,\n int in_features,\n int out_features,\n float eps) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx < num_groups) {\n int start = idx * in_features;\n int out_start = idx * out_features;\n \n float mean = 0.0f;\n for (int i = 0; i < in_features; i++) {\n mean += input[start + i];\n }\n mean /= static_cast<float>(in_features);\n \n float var = 0.0f;\n for (int i = 0; i < in_features; i++) {\n float diff = input[start + i] - mean;\n var += diff * diff;\n }\n var /= static_cast<float>(in_features);\n \n float inv_std = rsqrtf(var + eps);\n \n for (int o = 0; o < out_features; o++) {\n float q_sum = 0.0f, k_sum = 0.0f, v_sum = 0.0f;\n for (int i = 0; i < in_features; i++) {\n float normalized = (input[start + i] - mean) * inv_std * ln_weight[i] + ln_bias[i];\n q_sum += normalized * q_weight[o * in_features + i];\n k_sum += normalized * k_weight[o * in_features + i];\n v_sum += normalized * v_weight[o * in_features + i];\n }\n q_out[out_start + o] = q_sum;\n k_out[out_start + o] = k_sum;\n v_out[out_start + o] = v_sum;\n }\n }\n}\n\n__global__ void output_proj_kernel(\n const float* __restrict__ input,\n const float* __restrict__ weight,\n const float* __restrict__ bias,\n float* __restrict__ output,\n int num_groups,\n int in_features,\n int out_features) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx < num_groups) {\n int in_start = idx * in_features;\n int out_start = idx * out_features;\n \n for (int o = 0; o < out_features; o++) {\n float sum = bias[o];\n for (int i = 0; i < in_features; i++) {\n sum += input[in_start + i] * weight[o * in_features + i];\n }\n output[out_start + o] = sum;\n }\n }\n}\n\nstd::vector<torch::Tensor> msa_row_attn_cuda(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor q_weight,\n torch::Tensor k_weight,\n torch::Tensor v_weight,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n int num_heads,\n double eps) {\n \n TORCH_CHECK(msa.is_cuda(), \"MSA must be a CUDA tensor\");\n \n auto sizes = msa.sizes();\n int64_t B = sizes[0], S = sizes[1], N = sizes[2], D = sizes[3];\n int64_t d = D / num_heads;\n int64_t h = num_heads;\n \n int64_t num_groups = B * S * N;\n \n auto q_out = torch::empty({B, S, N, D}, msa.options());\n auto k_out = torch::empty({B, S, N, D}, msa.options());\n auto v_out = torch::empty({B, S, N, D}, msa.options());\n \n int blocks = (static_cast<int>(num_groups) + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n layernorm_multi_linear_kernel<<<blocks, BLOCK_SIZE>>>(\n msa.data_ptr<float>(),\n ln_weight.data_ptr<float>(),\n ln_bias.data_ptr<float>(),\n q_out.data_ptr<float>(),\n k_out.data_ptr<float>(),\n v_out.data_ptr<float>(),\n q_weight.data_ptr<float>(),\n k_weight.data_ptr<float>(),\n v_weight.data_ptr<float>(),\n static_cast<int>(num_groups),\n static_cast<int>(D),\n static_cast<int>(D),\n static_cast<float>(eps));\n \n // Reshape for attention using PyTorch operations (well-optimized)\n // [B, S, N, D] -> [B, S, N, h, d] -> [B, S, h, N, d]\n auto q_reshaped = q_out.view({B, S, N, h, d}).permute({0, 1, 3, 2, 4}).contiguous();\n auto k_reshaped = k_out.view({B, S, N, h, d}).permute({0, 1, 3, 2, 4}).contiguous();\n auto v_reshaped = v_out.view({B, S, N, h, d}).permute({0, 1, 3, 2, 4}).contiguous();\n \n // Compute attention: [B, S, h, N, d] @ [B, S, h, d, N] = [B, S, h, N, N]\n auto attn = torch::matmul(q_reshaped, k_reshaped.transpose(-1, -2)) / sqrtf(static_cast<float>(d));\n attn = torch::softmax(attn, -1);\n \n // Apply attention: [B, S, h, N, N] @ [B, S, h, N, d] = [B, S, h, N, d]\n auto attn_out = torch::matmul(attn, v_reshaped);\n \n // Reshape back: [B, S, h, N, d] -> [B, S, N, h, d] -> [B, S, N, D]\n auto attn_out_reshaped = attn_out.permute({0, 1, 3, 2, 4}).contiguous().view({B, S, N, D});\n \n // Output projection\n auto final_out = torch::empty({B, S, N, D}, msa.options());\n blocks = (static_cast<int>(num_groups) + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n output_proj_kernel<<<blocks, BLOCK_SIZE>>>(\n attn_out_reshaped.data_ptr<float>(),\n out_weight.data_ptr<float>(),\n out_bias.data_ptr<float>(),\n final_out.data_ptr<float>(),\n static_cast<int>(num_groups),\n static_cast<int>(D),\n static_cast<int>(D));\n \n return {final_out, q_out, k_out, v_out};\n}\n\n__global__ void outer_product_kernel(\n const float* __restrict__ input,\n const float* __restrict__ ln_weight,\n const float* __restrict__ ln_bias,\n const float* __restrict__ a_weight,\n const float* __restrict__ a_bias,\n const float* __restrict__ b_weight,\n const float* __restrict__ b_bias,\n float* __restrict__ a_out,\n float* __restrict__ b_out,\n int num_groups,\n int in_features,\n int hidden,\n float eps) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx < num_groups) {\n int in_start = idx * in_features;\n int out_start = idx * hidden;\n \n float mean = 0.0f;\n for (int i = 0; i < in_features; i++) {\n mean += input[in_start + i];\n }\n mean /= static_cast<float>(in_features);\n \n float var = 0.0f;\n for (int i = 0; i < in_features; i++) {\n float diff = input[in_start + i] - mean;\n var += diff * diff;\n }\n var /= static_cast<float>(in_features);\n \n float inv_std = rsqrtf(var + eps);\n \n for (int o = 0; o < hidden; o++) {\n float a_sum = a_bias[o];\n float b_sum = b_bias[o];\n for (int i = 0; i < in_features; i++) {\n float normalized = (input[in_start + i] - mean) * inv_std * ln_weight[i] + ln_bias[i];\n a_sum += normalized * a_weight[o * in_features + i];\n b_sum += normalized * b_weight[o * in_features + i];\n }\n a_out[out_start + o] = a_sum;\n b_out[out_start + o] = b_sum;\n }\n }\n}\n\ntorch::Tensor outer_product_mean_cuda(\n torch::Tensor msa,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor a_weight,\n torch::Tensor a_bias,\n torch::Tensor b_weight,\n torch::Tensor b_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n double eps) {\n \n TORCH_CHECK(msa.is_cuda(), \"MSA must be a CUDA tensor\");\n \n auto sizes = msa.sizes();\n int64_t B = sizes[0], S = sizes[1], N = sizes[2], D = sizes[3];\n int64_t hidden = a_weight.size(0);\n \n int64_t num_groups = B * S * N;\n \n auto a_out = torch::empty({B, S, N, hidden}, msa.options());\n auto b_out = torch::empty({B, S, N, hidden}, msa.options());\n \n int blocks = (static_cast<int>(num_groups) + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n outer_product_kernel<<<blocks, BLOCK_SIZE>>>(\n msa.data_ptr<float>(),\n ln_weight.data_ptr<float>(),\n ln_bias.data_ptr<float>(),\n a_weight.data_ptr<float>(),\n a_bias.data_ptr<float>(),\n b_weight.data_ptr<float>(),\n b_bias.data_ptr<float>(),\n a_out.data_ptr<float>(),\n b_out.data_ptr<float>(),\n static_cast<int>(num_groups),\n static_cast<int>(D),\n static_cast<int>(hidden),\n static_cast<float>(eps));\n \n auto outer = torch::einsum(\"bsih,bsjk->bijhk\", {a_out, b_out});\n outer = outer / static_cast<float>(S);\n \n int64_t N1 = outer.size(1), N2 = outer.size(2), H1 = outer.size(3), H2 = outer.size(4);\n auto outer_flat = outer.reshape({B, N1, N2, H1 * H2});\n \n return torch::linear(outer_flat, out_weight, out_bias);\n}\n\n__global__ void tri_mul_layernorm_linear_kernel(\n const float* __restrict__ input,\n const float* __restrict__ ln_weight,\n const float* __restrict__ ln_bias,\n const float* __restrict__ pa_weight,\n const float* __restrict__ pa_bias,\n const float* __restrict__ pb_weight,\n const float* __restrict__ pb_bias,\n const float* __restrict__ ga_weight,\n const float* __restrict__ ga_bias,\n const float* __restrict__ gb_weight,\n const float* __restrict__ gb_bias,\n float* __restrict__ pa_out,\n float* __restrict__ pb_out,\n int num_groups,\n int in_features,\n int hidden,\n float eps) {\n \n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (idx < num_groups) {\n int in_start = idx * in_features;\n int out_start = idx * hidden;\n \n float mean = 0.0f;\n for (int i = 0; i < in_features; i++) {\n mean += input[in_start + i];\n }\n mean /= static_cast<float>(in_features);\n \n float var = 0.0f;\n for (int i = 0; i < in_features; i++) {\n float diff = input[in_start + i] - mean;\n var += diff * diff;\n }\n var /= static_cast<float>(in_features);\n \n float inv_std = rsqrtf(var + eps);\n \n for (int o = 0; o < hidden; o++) {\n float pa_sum = pa_bias[o];\n float pb_sum = pb_bias[o];\n float ga_sum = ga_bias[o];\n float gb_sum = gb_bias[o];\n \n for (int i = 0; i < in_features; i++) {\n float normalized = (input[in_start + i] - mean) * inv_std * ln_weight[i] + ln_bias[i];\n pa_sum += normalized * pa_weight[o * in_features + i];\n pb_sum += normalized * pb_weight[o * in_features + i];\n ga_sum += normalized * ga_weight[o * in_features + i];\n gb_sum += normalized * gb_weight[o * in_features + i];\n }\n \n pa_out[out_start + o] = pa_sum * (1.0f / (1.0f + expf(-ga_sum)));\n pb_out[out_start + o] = pb_sum * (1.0f / (1.0f + expf(-gb_sum)));\n }\n }\n}\n\ntorch::Tensor tri_mul_out_cuda(\n torch::Tensor pair,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor pa_weight,\n torch::Tensor pa_bias,\n torch::Tensor pb_weight,\n torch::Tensor pb_bias,\n torch::Tensor ga_weight,\n torch::Tensor ga_bias,\n torch::Tensor gb_weight,\n torch::Tensor gb_bias,\n torch::Tensor norm_out_weight,\n torch::Tensor norm_out_bias,\n torch::Tensor out_weight,\n torch::Tensor out_bias,\n torch::Tensor gate_weight,\n torch::Tensor gate_bias,\n double eps) {\n \n TORCH_CHECK(pair.is_cuda(), \"Pair must be a CUDA tensor\");\n \n auto sizes = pair.sizes();\n int64_t B = sizes[0], N1 = sizes[1], N2 = sizes[2], D = sizes[3];\n int64_t hidden = pa_weight.size(0);\n \n int64_t num_groups = B * N1 * N2;\n \n auto pa_out = torch::empty({B, N1, N2, hidden}, pair.options());\n auto pb_out = torch::empty({B, N1, N2, hidden}, pair.options());\n \n int blocks = (static_cast<int>(num_groups) + BLOCK_SIZE - 1) / BLOCK_SIZE;\n \n tri_mul_layernorm_linear_kernel<<<blocks, BLOCK_SIZE>>>(\n pair.data_ptr<float>(),\n ln_weight.data_ptr<float>(),\n ln_bias.data_ptr<float>(),\n pa_weight.data_ptr<float>(),\n pa_bias.data_ptr<float>(),\n pb_weight.data_ptr<float>(),\n pb_bias.data_ptr<float>(),\n ga_weight.data_ptr<float>(),\n ga_bias.data_ptr<float>(),\n gb_weight.data_ptr<float>(),\n gb_bias.data_ptr<float>(),\n pa_out.data_ptr<float>(),\n pb_out.data_ptr<float>(),\n static_cast<int>(num_groups),\n static_cast<int>(D),\n static_cast<int>(hidden),\n static_cast<float>(eps));\n \n auto o = torch::einsum(\"bikh,bjkh->bijh\", {pa_out, pb_out});\n o = torch::layer_norm(o, {hidden}, norm_out_weight, norm_out_bias, eps);\n o = torch::linear(o, out_weight, out_bias);\n auto gate = torch::sigmoid(torch::linear(pair, gate_weight, gate_bias));\n \n return gate * o;\n}\n\"\"\"\n\nevoformer_ext = load_inline(\n name='evoformer_ext',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['msa_row_attn', 'outer_product_mean', 'tri_mul_out'],\n verbose=True\n)\n\n\nclass MSARowSelfAttention(nn.Module):\n def __init__(self, msa_dim, num_heads):\n super().__init__()\n self.num_heads = num_heads\n self.head_dim = msa_dim // num_heads\n self.norm = nn.LayerNorm(msa_dim)\n self.q = nn.Linear(msa_dim, msa_dim, bias=False)\n self.k = nn.Linear(msa_dim, msa_dim, bias=False)\n self.v = nn.Linear(msa_dim, msa_dim, bias=False)\n self.out = nn.Linear(msa_dim, msa_dim)\n\n def forward(self, msa):\n B, S, N, D = msa.shape\n \n results = evoformer_ext.msa_row_attn(\n msa.contiguous(),\n self.norm.weight,\n self.norm.bias,\n self.q.weight,\n self.k.weight,\n self.v.weight,\n self.out.weight,\n self.out.bias,\n self.num_heads,\n 1e-5\n )\n return results[0]\n\n\nclass OuterProductMeanSmall(nn.Module):\n def __init__(self, msa_dim, pair_dim, hidden=32):\n super().__init__()\n self.norm = nn.LayerNorm(msa_dim)\n self.a = nn.Linear(msa_dim, hidden)\n self.b = nn.Linear(msa_dim, hidden)\n self.out = nn.Linear(hidden * hidden, pair_dim)\n\n def forward(self, msa):\n return evoformer_ext.outer_product_mean(\n msa.contiguous(),\n self.norm.weight,\n self.norm.bias,\n self.a.weight,\n self.a.bias,\n self.b.weight,\n self.b.bias,\n self.out.weight,\n self.out.bias,\n 1e-5\n )\n\n\nclass TriMulOut(nn.Module):\n def __init__(self, pair_dim, hidden):\n super().__init__()\n self.norm = nn.LayerNorm(pair_dim)\n self.pa = nn.Linear(pair_dim, hidden)\n self.pb = nn.Linear(pair_dim, hidden)\n self.ga = nn.Linear(pair_dim, hidden)\n self.gb = nn.Linear(pair_dim, hidden)\n self.norm_out = nn.LayerNorm(hidden)\n self.out = nn.Linear(hidden, pair_dim)\n self.gate_out = nn.Linear(pair_dim, pair_dim)\n\n def forward(self, pair):\n return evoformer_ext.tri_mul_out(\n pair.contiguous(),\n self.norm.weight,\n self.norm.bias,\n self.pa.weight,\n self.pa.bias,\n self.pb.weight,\n self.pb.bias,\n self.ga.weight,\n self.ga.bias,\n self.gb.weight,\n self.gb.bias,\n self.norm_out.weight,\n self.norm_out.bias,\n self.out.weight,\n self.out.bias,\n self.gate_out.weight,\n self.gate_out.bias,\n 1e-5\n )\n\n\nclass ModelNew(nn.Module):\n def __init__(self, msa_dim, pair_dim, num_heads, tri_hidden):\n super().__init__()\n self.msa_attn = MSARowSelfAttention(msa_dim, num_heads)\n self.opm = OuterProductMeanSmall(msa_dim, pair_dim)\n self.tri_mul = TriMulOut(pair_dim, tri_hidden)\n\n def forward(\n self, msa_repr: torch.Tensor, pair_repr: torch.Tensor\n ) -> tuple:\n msa_repr = msa_repr + self.msa_attn(msa_repr)\n pair_repr = pair_repr + self.opm(msa_repr)\n pair_repr = pair_repr + self.tri_mul(pair_repr)\n return msa_repr, pair_repr"
}