def bdd_message_func(self, edges):
        """Message function for block-diagonal-decomposition regularizer"""
        ctx = edges.src["h"].context
        if (
            edges.src["h"].dtype in (np.int32, np.int64)
            and len(edges.src["h"].shape) == 1
        ):
            raise TypeError(
                "Block decomposition does not allow integer ID feature."
            )
        weight = self.weight.data(ctx)[edges.data["type"], :].reshape(
            -1, self.submat_in, self.submat_out
        )
        node = edges.src["h"].reshape(-1, 1, self.submat_in)
        msg = nd.batch_dot(node, weight).reshape(-1, self.out_feat)
        if "norm" in edges.data:
            msg = msg * edges.data["norm"]
        return {"msg": msg}
Message function for block-diagonal-decomposition regularizer
python/dgl/nn/mxnet/conv/relgraphconv.py