Project: s41467-024-49516-2 | Function: vector_rejection (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements vector rejection operation for maintaining equivariance in geometric vector features. Projects vectors onto directional edges and subtracts the projection to ensure rotational and translational invariance in protein structure representation.

**KEY TERMS:** {
  "vec": "vector features to be processed v",
  "d_ij": "normalized directional edge vector d_ij",
  "vec_proj": "vector projection (v · d_ij) onto edge direction",
  "return": "rejected vector maintaining equivariance"
}

**CONTEXT CODE:**         assert hidden_channels % num_heads == 0, (
            f"The number of hidden channels ({hidden_channels}) "
            f"must be evenly divisible by the number of "
            f"attention heads ({num_heads})"
        )

        self.distance_influence = distance_influence
        self.num_heads = num_heads
        self.hidden_channels = hidden_channels
        self.head_dim = hidden_channels // num_heads
        self.last_layer = last_layer
        self.vecnorm_type = vecnorm_type
        self.vecnorm_trainable = vecnorm_trainable

        self.x_layernorm = nn.LayerNorm(hidden_channels)  # 归一化x
        self.f_layernorm = nn.LayerNorm(hidden_channels)
        self.v_layernorm = VecLayerNorm(hidden_channels, vecnorm_trainable, vecnorm_type)
        self.act = activation()
        self.attn_activation = act_class_mapping[attn_activation]()
        self.cutoff = CosineCutoff(cutoff_lower, cutoff_upper)

        self.q_proj = nn.Linear(hidden_channels, hidden_channels)
        self.k_proj = nn.Linear(hidden_channels, hidden_channels)
        self.v_proj = nn.Linear(hidden_channels, hidden_channels)
        self.s_proj = nn.Linear(hidden_channels, hidden_channels * 2)

        self.v_dot_proj = nn.Linear(hidden_channels, hidden_channels)
        
        if not self.last_layer:
            self.f_proj = nn.Linear(hidden_channels, hidden_channels * 2)
            self.src_proj = nn.Linear(hidden_channels, hidden_channels, bias=False)
            self.trg_proj = nn.Linear(hidden_channels, hidden_channels, bias=False)
            self.w_dot_proj = nn.Linear(hidden_channels, hidden_channels)

        self.o_proj = nn.Linear(hidden_channels, hidden_channels * 3)
        self.vec_proj = nn.Linear(hidden_channels, hidden_channels * 3, bias=False)

        self.dk_proj = None
        if distance_influence in ["keys", "both"]:
            self.dk_proj = nn.Linear(hidden_channels, hidden_channels)

        self.dv_proj = None
        if distance_influence in ["values", "both"]:
            self.dv_proj = nn.Linear(hidden_channels, hidden_channels)
            
        self.scalar_dropout = nn.Dropout(dropout)
        self.vector_dropout = nn.Dropout2d(dropout)

        self.reset_parameters()
        
    def vector_rejection(self, vec, d_ij):
        """Implements vector rejection operation for maintaining equivariance in geometric vector features."""
        raise NotImplementedError("Function has been masked for testing purposes")

    def reset_parameters(self):
        self.x_layernorm.reset_parameters()
        self.f_layernorm.reset_parameters()
        self.v_layernorm.reset_parameters()

        nn.init.xavier_uniform_(self.q_proj.weight)
        self.q_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.k_proj.weight)
        self.k_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.v_proj.weight)
        self.v_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.o_proj.weight)
        self.o_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.s_proj.weight)
        self.s_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.v_dot_proj.weight)
        self.v_dot_proj.bias.data.fill_(0)
        
        if not self.last_layer:
            nn.init.xavier_uniform_(self.f_proj.weight)
            self.f_proj.bias.data.fill_(0)
            nn.init.xavier_uniform_(self.src_proj.weight)
            nn.init.xavier_uniform_(self.trg_proj.weight)
            nn.init.xavier_uniform_(self.w_dot_proj.weight)
            self.w_dot_proj.bias.data.fill_(0)

        nn.init.xavier_uniform_(self.vec_proj.weight)
        if self.dk_proj:
            nn.init.xavier_uniform_(self.dk_proj.weight)
            self.dk_proj.bias.data.fill_(0)
        if self.dv_proj:
            nn.init.xavier_uniform_(self.dv_proj.weight)
            self.dv_proj.bias.data.fill_(0)

    def forward(self, x, vec, edge_index, r_ij, f_ij, d_ij):
        #x, vec（点的向量表征）, edge_index, edge_weight（r_ij距离）, 
        # edge_attr（f_ij边属性）, edge_vec（d_ij边向量）
        # print()
        x = self.x_layernorm(x)
        f_ij = self.f_layernorm(f_ij)
        vec = self.v_layernorm(vec)
        # print("vec",vec.shape) 
        # vec不是edge_vec，最开始都是0，shape是x.size(0), 
        # ((self.lmax + 1) ** 2) - 1, hidden_channels,lmax=1时是3，lmax=2时是8
        
        q = self.q_proj(x).reshape(-1, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(-1, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(-1, self.num_heads, self.head_dim) 
        # 对x进行投影，然后reshape成num_heads个head_dim维的向量

        # print("q",q.shape) 
        # print("k",k.shape)
        # print("v",v.shape)

        vec1, vec2, vec3 = torch.split(self.vec_proj(vec), self.hidden_channels, dim=-1)
        # print(self.vec_proj(vec).shape) #把256通过linear变成768，然后分成3份，每份256，
        # 没有bias，意味着nn中只有乘法没有加法
        # print("vec1",vec1.shape) #vec1 shape是x.size(0), ((self.lmax + 1) ** 2) - 1, hidden_channels
        # print("vec2",vec2.shape) #vec2 shape是x.size(0), ((self.lmax + 1) ** 2) - 1, hidden_channels
        # print("vec3",vec3.shape) #vec3 shape是x.size(0), ((self.lmax + 1) ** 2) - 1, hidden_channels
        vec_dot = (vec1 * vec2).sum(dim=1)  #每个位置的元素相乘，然后在第一维上求和
        # print(vec1,vec2,vec1 * vec2,vec_dot)
        # print("vec_dot",vec_dot.shape) 38*256 if x.size(0)=38 and hidden_channels=256
        vec_dot = self.act(self.v_dot_proj(vec_dot))  # 过一个nn和激活函数

        dk = (
            self.act(self.dk_proj(f_ij)).reshape(-1, self.num_heads, self.head_dim)        #边属性投影1，转为和x一样的shape
            if self.dk_proj is not None
            else None
        )
        dv = (
            self.act(self.dv_proj(f_ij)).reshape(-1, self.num_heads, self.head_dim)          # 边属性投影2，转为和x一样的shape
            if self.dv_proj is not None
            else None
        )

        # propagate_type: (q: Tensor, k: Tensor, v: Tensor, vec: Tensor, dk: Tensor, dv: Tensor, r_ij: Tensor, d_ij: Tensor)
        x, vec_out = self.propagate(  # 在propagate中调用message和update，并不显式编辑propgate函数，
        # 运行时，qkv分别是每一条边的起点，终点的qkv，最终使用message和edge_update达到更新节点和边的效果
            edge_index,
            q=q,
            k=k,
            v=v,
            vec=vec,
            dk=dk,
            dv=dv,
            r_ij=r_ij,
            d_ij=d_ij,
            size=None,
        )
        # edge_updater_type: (vec: Tensor, d_ij: Tensor, f_ij: Tensor)
        o1, o2, o3 = torch.split(self.o_proj(x), self.hidden_channels, dim=1)  
        # output投影，分成3份，其中2份用来更新dx，1份用来更新dvec，dvec同时与vec_out也有关
        dx = vec_dot * o2 + o3
        dvec = vec3 * o1.unsqueeze(1) + vec_out
        if not self.last_layer:


Complete the function implementation replacing "raise NotImplementedError".