def bdd_message_func(self, edges):
        """Message function for block-diagonal-decomposition regularizer"""
        if (edges.src["h"].dtype == tf.int64) and len(
            edges.src["h"].shape
        ) == 1:
            raise TypeError(
                "Block decomposition does not allow integer ID feature."
            )

        # calculate msg @ W_r before put msg into edge
        # if src is th.int64 we expect it is an index select
        if self.low_mem:
            etypes, _ = tf.unique(edges.data["type"])
            msg = tf.zeros([edges.src["h"].shape[0], self.out_feat])
            idx = tf.range(edges.src["h"].shape[0])
            for etype in etypes:
                loc = edges.data["type"] == etype
                w = tf.reshape(
                    self.weight[etype],
                    (self.num_bases, self.submat_in, self.submat_out),
                )
                src = tf.reshape(
                    tf.boolean_mask(edges.src["h"], loc),
                    (-1, self.num_bases, self.submat_in),
                )
                sub_msg = tf.einsum("abc,bcd->abd", src, w)
                sub_msg = tf.reshape(sub_msg, (-1, self.out_feat))
                indices = tf.reshape(tf.boolean_mask(idx, loc), (-1, 1))
                msg = tf.tensor_scatter_nd_update(msg, indices, sub_msg)
        else:
            weight = tf.reshape(
                tf.gather(self.weight, edges.data["type"]),
                (-1, self.submat_in, self.submat_out),
            )
            node = tf.reshape(edges.src["h"], (-1, 1, self.submat_in))
            msg = tf.reshape(tf.matmul(node, weight), (-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/tensorflow/conv/relgraphconv.py