def forward(self, x, x_type, sorted_by_type=False):
        """Forward computation.

        Parameters
        ----------
        x : torch.Tensor
            A 2D input tensor. Shape: (N, D1)
        x_type : torch.Tensor
            A 1D integer tensor storing the type of the elements in ``x`` with one-to-one
            correspondenc. Shape: (N,)
        sorted_by_type : bool, optional
            Whether the inputs have been sorted by the types. Forward on pre-sorted inputs may
            be faster.

        Returns
        -------
        y : torch.Tensor
            The transformed output tensor. Shape: (N, D2)
        """
        w = self.get_weight()
        if self.regularizer == "bdd":
            w = w.index_select(0, x_type).view(
                -1, self.submat_in, self.submat_out
            )
            x = x.view(-1, 1, self.submat_in)
            return torch.bmm(x, w).view(-1, self.out_size)
        elif sorted_by_type:
            pos_l = torch.searchsorted(
                x_type, torch.arange(self.num_types, device=x.device)
            )
            pos_r = torch.cat(
                [pos_l[1:], torch.tensor([len(x_type)], device=x.device)]
            )
            seglen = (
                pos_r - pos_l
            ).cpu()  # XXX(minjie): cause device synchronize
            return segment_mm(x, w, seglen_a=seglen)
        else:
            return gather_mm(x, w, idx_b=x_type)
Forward computation.

        Parameters
        ----------
        x : torch.Tensor
            A 2D input tensor. Shape: (N, D1)
        x_type : torch.Tensor
            A 1D integer tensor storing the type of the elements in ``x`` with one-to-one
            correspondenc. Shape: (N,)
        sorted_by_type : bool, optional
            Whether the inputs have been sorted by the types. Forward on pre-sorted inputs may
            be faster.

        Returns
        -------
        y : torch.Tensor
            The transformed output tensor. Shape: (N, D2)
        
python/dgl/nn/pytorch/linear.py