Skip to content

Hetero API

Generated reference for heterogeneous graphs — multiple node types and relations (typed edges). Start with the overview (academic author / paper / institution walkthrough) and examples.

Model zoo acronyms (R-GCN, GraphSAGE, HAN, HGT, CompGCN) are spelled out in the model zoo table.

anytensor.hetero

Heterogeneous graph types and tree batch/unbatch hooks.

Not part of the jraph-mirroring API. Import from here::

from anytensor.hetero import HeteroGraphsTuple, multi_update_all
from anytensor.hetero import relational_graph_convolution

HeteroGraphsTuple

Bases: NamedTuple

Heterogeneous graph(s) with per-type node pools and per-relation incidence.

Node ids in senders / receivers for etype (src, rel, dst) are local to nodes[src] and nodes[dst] (not a global node pool).

n_node[ntype] and n_edge[etype] are integer vectors of length n_graphs (jraph-style batching within one object).

:func:anytensor.tree.batch requires every input to share the same keys.

Source code in anytensor/hetero/graph.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
class HeteroGraphsTuple(NamedTuple):
    """Heterogeneous graph(s) with per-type node pools and per-relation incidence.

    Node ids in ``senders`` / ``receivers`` for etype ``(src, rel, dst)`` are
    local to ``nodes[src]`` and ``nodes[dst]`` (not a global node pool).

    ``n_node[ntype]`` and ``n_edge[etype]`` are integer vectors of length
    ``n_graphs`` (jraph-style batching within one object).

    :func:`anytensor.tree.batch` requires every input to share the same keys.
    """

    nodes: Mapping[Ntype, Optional[ArrayTree]]
    edges: Mapping[CanonicalEtype, Optional[ArrayTree]]
    senders: Mapping[CanonicalEtype, Any]
    receivers: Mapping[CanonicalEtype, Any]
    n_node: Mapping[Ntype, Any]
    n_edge: Mapping[CanonicalEtype, Any]
    globals: Optional[ArrayTree] = None

    def n_graphs(self) -> int:
        return _n_graphs_from_sizes(self.n_node, self.globals)

    def ntypes(self) -> Tuple[str, ...]:
        return tuple(sorted(self.n_node.keys()))

    def canonical_etypes(self) -> Tuple[CanonicalEtype, ...]:
        return tuple(sorted(self.n_edge.keys()))

    def update(
        self,
        nodes: Optional[Mapping[Ntype, Optional[ArrayTree]]] = None,
        edges: Optional[Mapping[CanonicalEtype, Optional[ArrayTree]]] = None,
        senders: Optional[Mapping[CanonicalEtype, Any]] = None,
        receivers: Optional[Mapping[CanonicalEtype, Any]] = None,
        n_node: Optional[Mapping[Ntype, Any]] = None,
        n_edge: Optional[Mapping[CanonicalEtype, Any]] = None,
        globals: Any = _UNSET,
    ) -> "HeteroGraphsTuple":
        """Return a new graph with shallow-merged mapping fields.

        Only keys present in the update dicts are replaced; other keys are
        kept. Pass ``globals=...`` to replace globals (including with ``None``).
        """
        return HeteroGraphsTuple(
            nodes=_merge_map(self.nodes, nodes),
            edges=_merge_map(self.edges, edges),
            senders=_merge_map(self.senders, senders),
            receivers=_merge_map(self.receivers, receivers),
            n_node=_merge_map(self.n_node, n_node),
            n_edge=_merge_map(self.n_edge, n_edge),
            globals=self.globals if globals is _UNSET else globals,
        )

    def iter_nodes(
        self, *, skip_empty: bool = False
    ) -> Iterator[Tuple[Ntype, Optional[ArrayTree], Any]]:
        for ntype in self.ntypes():
            n = self.n_node[ntype]
            if skip_empty and _sum_int(n) == 0:
                continue
            yield ntype, self.nodes.get(ntype), n

    def iter_edges(
        self, *, skip_empty: bool = False
    ) -> Iterator[Tuple[CanonicalEtype, Optional[ArrayTree], Any]]:
        for etype in self.canonical_etypes():
            n = self.n_edge[etype]
            if skip_empty and _sum_int(n) == 0:
                continue
            yield etype, self.edges.get(etype), n

    def iter_relations(
        self,
        etypes: Optional[Sequence[CanonicalEtype]] = None,
        *,
        skip_empty: bool = True,
        reverse: bool = False,
    ) -> Iterator[SendRecvTuple]:
        """Yield aliasing :class:`SendRecvTuple` views for each relation."""
        keys = list(etypes) if etypes is not None else list(self.canonical_etypes())
        for etype in keys:
            if etype not in self.n_edge:
                continue
            if skip_empty and _sum_int(self.n_edge[etype]) == 0:
                continue
            yield self.relation_view(etype)
            if reverse:
                yield self.relation_view(etype, reverse=True)

    def relation_view(
        self, etype: CanonicalEtype, *, reverse: bool = False
    ) -> SendRecvTuple:
        """Aliasing send→recv view for one canonical etype."""
        src, _rel, dst = etype
        if reverse:
            return SendRecvTuple(
                nodes_send=self.nodes.get(dst),
                nodes_recv=self.nodes.get(src),
                senders=self.receivers[etype],
                receivers=self.senders[etype],
                edges=self.edges.get(etype),
                n_node_send=self.n_node[dst],
                n_node_recv=self.n_node[src],
                n_edge=self.n_edge[etype],
                globals=self.globals,
                src_ntype=dst,
                dst_ntype=src,
                etype=(dst, _rel, src),
            )
        return SendRecvTuple(
            nodes_send=self.nodes.get(src),
            nodes_recv=self.nodes.get(dst),
            senders=self.senders[etype],
            receivers=self.receivers[etype],
            edges=self.edges.get(etype),
            n_node_send=self.n_node[src],
            n_node_recv=self.n_node[dst],
            n_edge=self.n_edge[etype],
            globals=self.globals,
            src_ntype=src,
            dst_ntype=dst,
            etype=etype,
        )

    @classmethod
    def __tree_batch__(cls, xs, axis: int = 0):
        if axis != 0:
            raise ValueError("HeteroGraphsTuple batch only supports axis=0")
        return _batch_hetero(xs)

    def __tree_unbatch__(self, axis: int = 0):
        if axis != 0:
            raise ValueError("HeteroGraphsTuple unbatch only supports axis=0")
        return _unbatch_hetero(self)

update 

update(
    nodes=None,
    edges=None,
    senders=None,
    receivers=None,
    n_node=None,
    n_edge=None,
    globals=_UNSET,
)

Return a new graph with shallow-merged mapping fields.

Only keys present in the update dicts are replaced; other keys are kept. Pass globals=... to replace globals (including with None).

Source code in anytensor/hetero/graph.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def update(
    self,
    nodes: Optional[Mapping[Ntype, Optional[ArrayTree]]] = None,
    edges: Optional[Mapping[CanonicalEtype, Optional[ArrayTree]]] = None,
    senders: Optional[Mapping[CanonicalEtype, Any]] = None,
    receivers: Optional[Mapping[CanonicalEtype, Any]] = None,
    n_node: Optional[Mapping[Ntype, Any]] = None,
    n_edge: Optional[Mapping[CanonicalEtype, Any]] = None,
    globals: Any = _UNSET,
) -> "HeteroGraphsTuple":
    """Return a new graph with shallow-merged mapping fields.

    Only keys present in the update dicts are replaced; other keys are
    kept. Pass ``globals=...`` to replace globals (including with ``None``).
    """
    return HeteroGraphsTuple(
        nodes=_merge_map(self.nodes, nodes),
        edges=_merge_map(self.edges, edges),
        senders=_merge_map(self.senders, senders),
        receivers=_merge_map(self.receivers, receivers),
        n_node=_merge_map(self.n_node, n_node),
        n_edge=_merge_map(self.n_edge, n_edge),
        globals=self.globals if globals is _UNSET else globals,
    )

iter_relations 

iter_relations(
    etypes=None, *, skip_empty=True, reverse=False
)

Yield aliasing :class:SendRecvTuple views for each relation.

Source code in anytensor/hetero/graph.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def iter_relations(
    self,
    etypes: Optional[Sequence[CanonicalEtype]] = None,
    *,
    skip_empty: bool = True,
    reverse: bool = False,
) -> Iterator[SendRecvTuple]:
    """Yield aliasing :class:`SendRecvTuple` views for each relation."""
    keys = list(etypes) if etypes is not None else list(self.canonical_etypes())
    for etype in keys:
        if etype not in self.n_edge:
            continue
        if skip_empty and _sum_int(self.n_edge[etype]) == 0:
            continue
        yield self.relation_view(etype)
        if reverse:
            yield self.relation_view(etype, reverse=True)

relation_view 

relation_view(etype, *, reverse=False)

Aliasing send→recv view for one canonical etype.

Source code in anytensor/hetero/graph.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def relation_view(
    self, etype: CanonicalEtype, *, reverse: bool = False
) -> SendRecvTuple:
    """Aliasing send→recv view for one canonical etype."""
    src, _rel, dst = etype
    if reverse:
        return SendRecvTuple(
            nodes_send=self.nodes.get(dst),
            nodes_recv=self.nodes.get(src),
            senders=self.receivers[etype],
            receivers=self.senders[etype],
            edges=self.edges.get(etype),
            n_node_send=self.n_node[dst],
            n_node_recv=self.n_node[src],
            n_edge=self.n_edge[etype],
            globals=self.globals,
            src_ntype=dst,
            dst_ntype=src,
            etype=(dst, _rel, src),
        )
    return SendRecvTuple(
        nodes_send=self.nodes.get(src),
        nodes_recv=self.nodes.get(dst),
        senders=self.senders[etype],
        receivers=self.receivers[etype],
        edges=self.edges.get(etype),
        n_node_send=self.n_node[src],
        n_node_recv=self.n_node[dst],
        n_edge=self.n_edge[etype],
        globals=self.globals,
        src_ntype=src,
        dst_ntype=dst,
        etype=etype,
    )

SendRecvTuple

Bases: NamedTuple

One-way send→receive incidence (homo or bipartite).

Views alias parent storage: nodes_send / nodes_recv may be the same object (homo) or two ntype pools (hetero relation).

Source code in anytensor/hetero/graph.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class SendRecvTuple(NamedTuple):
    """One-way send→receive incidence (homo or bipartite).

    Views alias parent storage: ``nodes_send`` / ``nodes_recv`` may be the
    same object (homo) or two ntype pools (hetero relation).
    """

    nodes_send: Optional[ArrayTree]
    nodes_recv: Optional[ArrayTree]
    senders: Any
    receivers: Any
    edges: Optional[ArrayTree]
    n_node_send: Any
    n_node_recv: Any
    n_edge: Any
    globals: Optional[ArrayTree] = None
    src_ntype: Optional[str] = None
    dst_ntype: Optional[str] = None
    etype: Optional[CanonicalEtype] = None

RelationSpec

Bases: NamedTuple

Per-relation update for :func:multi_update_all.

Attributes:

Name Type Description
message_fn MessageFn

(src, dst, edges) -> messages.
reduce ReduceName

Segment reduce after optional attention (sum with attention).
attention_logit_fn Optional[AttentionLogitFn]

Optional (src, dst, edges) -> logits.
attention_reduce_fn Optional[AttentionReduceFn]

Optional (messages, weights) -> messages. Defaults to element-wise multiply when only attention_logit_fn is set.
Source code in anytensor/hetero/message.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class RelationSpec(NamedTuple):
    """Per-relation update for :func:`multi_update_all`.

    Attributes:
        message_fn: ``(src, dst, edges) -> messages``.
        reduce: Segment reduce after optional attention (``sum`` with attention).
        attention_logit_fn: Optional ``(src, dst, edges) -> logits``.
        attention_reduce_fn: Optional ``(messages, weights) -> messages``.
            Defaults to element-wise multiply when only ``attention_logit_fn``
            is set.
    """

    message_fn: MessageFn
    reduce: ReduceName = "sum"
    attention_logit_fn: Optional[AttentionLogitFn] = None
    attention_reduce_fn: Optional[AttentionReduceFn] = None

graphs_tuple_as_send_recv 

graphs_tuple_as_send_recv(graph)

View a jraph :class:~anytensor.jraph.GraphsTuple as send→recv (aliased pools).

Source code in anytensor/hetero/graph.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def graphs_tuple_as_send_recv(graph) -> SendRecvTuple:
    """View a jraph :class:`~anytensor.jraph.GraphsTuple` as send→recv (aliased pools)."""
    return SendRecvTuple(
        nodes_send=graph.nodes,
        nodes_recv=graph.nodes,
        senders=graph.senders,
        receivers=graph.receivers,
        edges=graph.edges,
        n_node_send=graph.n_node,
        n_node_recv=graph.n_node,
        n_edge=graph.n_edge,
        globals=graph.globals,
        src_ntype=None,
        dst_ntype=None,
        etype=None,
    )

key_schema 

key_schema(g)

Sorted (ntypes, etypes) from present keys (empties still count).

Source code in anytensor/hetero/graph.py 
58
59
60
def key_schema(g: "HeteroGraphsTuple") -> Tuple[Tuple[str, ...], Tuple[CanonicalEtype, ...]]:
    """Sorted ``(ntypes, etypes)`` from present keys (empties still count)."""
    return tuple(sorted(g.n_node.keys())), tuple(sorted(g.n_edge.keys()))

schemas_equal 

schemas_equal(a, b)

True if a and b have the same ntype / etype key sets.

Source code in anytensor/hetero/graph.py 
63
64
65
def schemas_equal(a: "HeteroGraphsTuple", b: "HeteroGraphsTuple") -> bool:
    """True if ``a`` and ``b`` have the same ntype / etype key sets."""
    return key_schema(a) == key_schema(b)

attention_weight_messages 

attention_weight_messages(messages, weights)

Default attention reduce: element-wise messages * weights.

Same pattern as Graph Attention Networks (GAT): after :func:~anytensor.segment.segment_softmax, multiply messages by the per-edge weights.

Source code in anytensor/hetero/message.py 
91
92
93
94
95
96
97
98
def attention_weight_messages(messages: ArrayTree, weights: ArrayTree) -> ArrayTree:
    """Default attention reduce: element-wise ``messages * weights``.

    Same pattern as Graph Attention Networks (GAT): after
    :func:`~anytensor.segment.segment_softmax`, multiply messages by the
    per-edge weights.
    """
    return tree.map(lambda m, w: m * w, messages, weights)

copy_u_message 

copy_u_message(src_nodes, dst_nodes, edges)

DGL fn.copy_u: message is the source node feature (ignore dst/edge).

Source code in anytensor/hetero/message.py 
115
116
117
118
def copy_u_message(src_nodes, dst_nodes, edges):
    """DGL ``fn.copy_u``: message is the source node feature (ignore dst/edge)."""
    del dst_nodes, edges
    return src_nodes

multi_update_all 

multi_update_all(
    graph,
    etype_dict=None,
    cross_reducer="sum",
    *,
    reduce="sum",
    etypes=None,
)

Multi-relation update aligned with DGL multi_update_all.

Per etype: message + optional attention + segment-reduce onto destination nodes. Then fuse mailboxes that share a destination ntype with cross_reducer.

Parameters:

Name Type Description Default
graph HeteroGraphsTuple

Heterogeneous graph(s).

 required
etype_dict Optional[Mapping[CanonicalEtype, Union[MessageFn, RelationSpec, tuple]]]

Optional map etype -> :class:RelationSpec, message_fn, (message_fn, reduce), or (message_fn, reduce, attention_logit_fn, attention_reduce_fn). Default: copy_u + reduce for every etype in etypes / canonical_etypes().

 None
cross_reducer CrossReduceName

Fuse per-relation mailboxes for the same destination ntype. sum / mean / max / min / stack match DGL when per-relation mailboxes match (see module docstring). stack uses axis 1 (DGL shape (n_dst, n_relations, ...)); order is etype_dict insertion order.

 'sum'
reduce ReduceName

Default per-relation segment reduce when not set in etype_dict. sum/mean/max/min match DGL (empty destinations 0).

 'sum'
etypes Optional[Sequence[CanonicalEtype]]

Subset of relations when etype_dict is omitted.

 None

Apply is @cache (same pattern as GraphNetwork).

Returns:

Type Description
HeteroGraphsTuple

A new :class:HeteroGraphsTuple whose destination node features are
HeteroGraphsTuple

replaced by the cross-reduced mailboxes (same as DGL writing the
HeteroGraphsTuple

reduced feature). Source-only ntypes are unchanged.
Source code in anytensor/hetero/message.py 
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
@cache
def multi_update_all(
    graph: HeteroGraphsTuple,
    etype_dict: Optional[
        Mapping[CanonicalEtype, Union[MessageFn, RelationSpec, tuple]]
    ] = None,
    cross_reducer: CrossReduceName = "sum",
    *,
    reduce: ReduceName = "sum",
    etypes: Optional[Sequence[CanonicalEtype]] = None,
) -> HeteroGraphsTuple:
    """Multi-relation update aligned with DGL ``multi_update_all``.

    Per etype: message + optional attention + segment-reduce onto destination
    nodes. Then fuse mailboxes that share a destination ntype with
    ``cross_reducer``.

    Args:
        graph: Heterogeneous graph(s).
        etype_dict: Optional map ``etype ->`` :class:`RelationSpec`,
            ``message_fn``, ``(message_fn, reduce)``, or
            ``(message_fn, reduce, attention_logit_fn, attention_reduce_fn)``.
            Default: ``copy_u`` + ``reduce`` for every etype in ``etypes`` /
            ``canonical_etypes()``.
        cross_reducer: Fuse per-relation mailboxes for the same destination
            ntype. ``sum`` / ``mean`` / ``max`` / ``min`` / ``stack`` match DGL
            when per-relation mailboxes match (see module docstring). ``stack``
            uses axis ``1`` (DGL shape ``(n_dst, n_relations, ...)``); order is
            ``etype_dict`` insertion order.
        reduce: Default per-relation segment reduce when not set in
            ``etype_dict``. ``sum``/``mean``/``max``/``min`` match DGL
            (empty destinations ``0``).
        etypes: Subset of relations when ``etype_dict`` is omitted.

    Apply is ``@cache`` (same pattern as GraphNetwork).

    Returns:
        A new :class:`HeteroGraphsTuple` whose destination node features are
        replaced by the cross-reduced mailboxes (same as DGL writing the
        reduced feature). Source-only ntypes are unchanged.
    """
    if etype_dict is None:
        keys = list(etypes) if etypes is not None else list(graph.canonical_etypes())
        etype_dict = {e: copy_u_message for e in keys}
    if not etype_dict:
        return graph

    mailboxes: dict[CanonicalEtype, ArrayTree] = {}
    for etype, raw in etype_dict.items():
        spec = _parse_relation_spec(raw, reduce)
        mailboxes[etype] = relation_mailbox(
            graph,
            etype,
            message_fn=spec.message_fn,
            reduce=spec.reduce,
            attention_logit_fn=spec.attention_logit_fn,
            attention_reduce_fn=spec.attention_reduce_fn,
        )

    new_nodes = dict(graph.nodes)
    dst_types = {e[2] for e in mailboxes}
    for ntype in dst_types:
        parts = [mailboxes[e] for e in mailboxes if e[2] == ntype]
        new_nodes[ntype] = _cross_reduce(parts, cross_reducer)
    return graph.update(nodes=new_nodes)

relation_mailbox 

relation_mailbox(
    graph,
    etype,
    *,
    message_fn=copy_u_message,
    reduce="sum",
    attention_logit_fn=None,
    attention_reduce_fn=None,
)

Per-relation messages reduced onto destination nodes (DGL type-wise step).

With :func:copy_u_message, reduce matches DGL fn.copy_u + fn.sum/fn.mean/fn.max/fn.min. Empty destinations are 0 (including max/min via segment_*_or_constant).

When attention_logit_fn is set, logits are softmax-normalized per destination (receivers) and attention_reduce_fn weights messages before the segment reduce — same flow as :func:anytensor.jraph.GraphNetwork attention. Omit attention_reduce_fn to default to :func:attention_weight_messages. With attention, prefer reduce="sum". Apply is @cache (same pattern as GraphNetwork).

Source code in anytensor/hetero/message.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@cache
def relation_mailbox(
    graph: HeteroGraphsTuple,
    etype: CanonicalEtype,
    *,
    message_fn: MessageFn = copy_u_message,
    reduce: ReduceName = "sum",
    attention_logit_fn: Optional[AttentionLogitFn] = None,
    attention_reduce_fn: Optional[AttentionReduceFn] = None,
):
    """Per-relation messages reduced onto destination nodes (DGL type-wise step).

    With :func:`copy_u_message`, ``reduce`` matches DGL ``fn.copy_u`` +
    ``fn.sum``/``fn.mean``/``fn.max``/``fn.min``. Empty destinations are ``0``
    (including max/min via ``segment_*_or_constant``).

    When ``attention_logit_fn`` is set, logits are softmax-normalized per
    destination (``receivers``) and ``attention_reduce_fn`` weights messages
    before the segment reduce — same flow as
    :func:`anytensor.jraph.GraphNetwork` attention. Omit
    ``attention_reduce_fn`` to default to :func:`attention_weight_messages`.
    With attention, prefer ``reduce="sum"``. Apply is ``@cache`` (same
    pattern as GraphNetwork).
    """
    if etype not in graph.n_edge:
        raise KeyError(f"etype {etype!r} not in graph")
    if reduce not in _SEGMENT_REDUCE:
        raise ValueError(f"unknown reduce {reduce!r}")
    if attention_reduce_fn is not None and attention_logit_fn is None:
        raise ValueError("attention_logit_fn is required when attention_reduce_fn is set")
    if attention_logit_fn is not None and attention_reduce_fn is None:
        attention_reduce_fn = attention_weight_messages

    src, _rel, dst = etype
    receivers = graph.receivers[etype]
    src_nodes = _take_nodes(graph.nodes[src], graph.senders[etype])
    dst_nodes = _take_nodes(graph.nodes[dst], receivers)
    edges = graph.edges.get(etype)
    messages = message_fn(src_nodes, dst_nodes, edges)
    num_dst = _leading(graph.nodes[dst])

    if attention_logit_fn is not None:
        logits = attention_logit_fn(src_nodes, dst_nodes, edges)
        weights = tree.map(
            lambda logit: segment_softmax(logit, receivers, num_dst),
            logits,
        )
        messages = attention_reduce_fn(messages, weights)

    return _reduce_messages(messages, receivers, num_dst, reduce)

comp_gcn 

comp_gcn(
    graph,
    relation_apply,
    self_apply,
    *,
    composition="mult",
    activation=_relu,
    reducer="sum",
)

CompGCN (Composition-based Multi-Relational GCN) layer.

Vashishth et al., ICLR 2020. Requires edge features on each used etype.

  • mult — relation_apply[r](h_src * e)
  • sum — relation_apply[r](h_src + e)

Then segment reducer, cross sum, and activation(self_apply[n](h) + mailbox).

Source code in anytensor/hetero/models.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
@cache
def comp_gcn(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    self_apply: Mapping[Ntype, LinearFn],
    *,
    composition: str = "mult",
    activation: ActivationFn = _relu,
    reducer: str = "sum",
) -> HeteroGraphsTuple:
    """CompGCN (Composition-based Multi-Relational GCN) layer.

    Vashishth et al., ICLR 2020. Requires edge features on each used etype.

    * ``mult`` — ``relation_apply[r](h_src * e)``
    * ``sum`` — ``relation_apply[r](h_src + e)``

    Then segment ``reducer``, cross ``sum``, and
    ``activation(self_apply[n](h) + mailbox)``.
    """
    if composition not in ("mult", "sum"):
        raise ValueError("composition must be 'mult' or 'sum'")

    def make_msg(apply: LinearFn):
        def msg(src, dst, edges):
            del dst
            if edges is None:
                raise ValueError("comp_gcn requires edge features on every etype")
            composed = src * edges if composition == "mult" else src + edges
            return apply(composed)

        return msg

    etype_dict = {
        etype: RelationSpec(message_fn=make_msg(fn), reduce=reducer)  # type: ignore[arg-type]
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(self_apply[ntype](graph.nodes[ntype]) + mail.nodes[ntype])
        for ntype in self_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

gat_attention_logit 

gat_attention_logit(
    src, dst, attn_vec_apply, *, negative_slope=0.2
)

GAT-style edge score: LeakyReLU(attn_vec_apply(concat(src, dst))).

Graph Attention Network (GAT; Veličković et al., ICLR 2018) neighborhood scoring. attn_vec_apply maps concatenated features to shape (E, 1) or (E,). Typical use inside a logit callable::

lambda s, d, e: gat_attention_logit(s, d, my_linear)
Source code in anytensor/hetero/models.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def gat_attention_logit(
    src: ArrayTree,
    dst: ArrayTree,
    attn_vec_apply: LinearFn,
    *,
    negative_slope: float = 0.2,
) -> ArrayTree:
    """GAT-style edge score: ``LeakyReLU(attn_vec_apply(concat(src, dst)))``.

    Graph Attention Network (GAT; Veličković et al., ICLR 2018) neighborhood
    scoring. ``attn_vec_apply`` maps concatenated features to shape ``(E, 1)``
    or ``(E,)``. Typical use inside a logit callable::

        lambda s, d, e: gat_attention_logit(s, d, my_linear)
    """
    cat = concatenate([src, dst], axis=-1)
    x = attn_vec_apply(cat)
    return where(x > 0, x, x * negative_slope)

han 

han(
    graph,
    meta_path_etypes,
    node_message,
    node_attention_logit,
    semantic_project,
    semantic_query,
    *,
    node_activation=_relu,
    semantic_activation=_tanh,
)

HAN (Heterogeneous Graph Attention Network) layer.

Wang et al., WWW 2019: node-level + semantic attention. Each meta_path_etypes entry is a meta-path hop already stored as a canonical etype (precompute longer paths as their own etypes).

  1. Node-level attention — node_message[e](h_src), logits from node_attention_logit[e](src, dst, edges), softmax over neighbors (GAT-style).
  2. Mailboxes stacked; semantic attention mixes path embeddings with semantic_query after semantic_project. Apply is @cache (same pattern as GraphNetwork).
Source code in anytensor/hetero/models.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
@cache
def han(
    graph: HeteroGraphsTuple,
    meta_path_etypes: Sequence[CanonicalEtype],
    node_message: Mapping[CanonicalEtype, LinearFn],
    node_attention_logit: Mapping[CanonicalEtype, LogitFn],
    semantic_project: LinearFn,
    semantic_query: ArrayTree,
    *,
    node_activation: ActivationFn = _relu,
    semantic_activation: ActivationFn = _tanh,
) -> HeteroGraphsTuple:
    """HAN (Heterogeneous Graph Attention Network) layer.

    Wang et al., WWW 2019: node-level + semantic attention. Each
    ``meta_path_etypes`` entry is a meta-path hop already stored as a
    canonical etype (precompute longer paths as their own etypes).

    1. **Node-level attention** — ``node_message[e](h_src)``, logits from
       ``node_attention_logit[e](src, dst, edges)``, softmax over neighbors
       (GAT-style).
    2. Mailboxes **stacked**; **semantic attention** mixes path embeddings
       with ``semantic_query`` after ``semantic_project``.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """
    if not meta_path_etypes:
        raise ValueError("han requires at least one meta-path etype")

    etype_dict = {
        etype: RelationSpec(
            message_fn=_src_message(node_message[etype]),
            reduce="sum",
            attention_logit_fn=node_attention_logit[etype],
            attention_reduce_fn=attention_weight_messages,
        )
        for etype in meta_path_etypes
    }
    stacked = multi_update_all(graph, etype_dict, cross_reducer="stack")
    nodes = dict(graph.nodes)
    for ntype in {e[2] for e in meta_path_etypes}:
        h_stack = stacked.nodes[ntype]  # (n, R, d)
        h_act = node_activation(h_stack)
        n, r, d = shape(h_act)[:3]
        flat = reshape(h_act, (-1, d))
        proj = semantic_project(flat)
        d_s = shape(proj)[-1]
        proj = reshape(proj, (n, r, d_s))
        proj = semantic_activation(proj)
        q_vec = reshape(semantic_query, (d_s,))
        score = at_sum(proj * q_vec, axes=-1)
        alpha = reshape(_softmax_axis1(score), (n, r, 1))
        nodes[ntype] = at_sum(h_act * alpha, axes=1)
    return graph.update(nodes=nodes)

hetero_sage 

hetero_sage(
    graph,
    relation_apply,
    combine_apply,
    *,
    activation=_relu,
)

Heterogeneous GraphSAGE mean layer (Hamilton et al., NeurIPS 2017).

GraphSAGE (SAmple and aggreGatE): per-relation map on sources, mean aggregate, cross sum, then activation(combine_apply[n](concat[h_self, mailbox])). Apply is @cache (same pattern as GraphNetwork).

Source code in anytensor/hetero/models.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@cache
def hetero_sage(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    combine_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _relu,
) -> HeteroGraphsTuple:
    """Heterogeneous GraphSAGE mean layer (Hamilton et al., NeurIPS 2017).

    GraphSAGE (SAmple and aggreGatE): per-relation map on sources, ``mean``
    aggregate, cross ``sum``, then
    ``activation(combine_apply[n](concat[h_self, mailbox]))``.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """
    etype_dict = {
        etype: RelationSpec(message_fn=_src_message(fn), reduce="mean")
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(
            combine(concatenate([graph.nodes[ntype], mail.nodes[ntype]], axis=-1))
        )
        for ntype, combine in combine_apply.items()
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

hgt 

hgt(
    graph,
    message_apply,
    attention_logit,
    target_apply,
    *,
    activation=_identity,
    scale=None,
)

HGT (Heterogeneous Graph Transformer) style layer.

Hu et al., WWW 2020 — typed attention + target projection. Full HGT uses typed query/key/value and edge-type matrices (often multi-head). Fold those into the callables you pass:

  • message_apply[etype](h_src) — value / message projection.
  • attention_logit[etype](src, dst, edges) — edge logits (include 1/sqrt(d) here, or set scale).
  • Softmax over neighbors, weighted sum, cross sum across etypes.
  • target_apply[ntype] — target-type output projection. Apply is @cache (same pattern as GraphNetwork).
Source code in anytensor/hetero/models.py 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
@cache
def hgt(
    graph: HeteroGraphsTuple,
    message_apply: Mapping[CanonicalEtype, LinearFn],
    attention_logit: Mapping[CanonicalEtype, LogitFn],
    target_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _identity,
    scale: Optional[float] = None,
) -> HeteroGraphsTuple:
    """HGT (Heterogeneous Graph Transformer) style layer.

    Hu et al., WWW 2020 — typed attention + target projection. Full HGT uses
    typed query/key/value and edge-type matrices (often multi-head). Fold
    those into the callables you pass:

    * ``message_apply[etype](h_src)`` — value / message projection.
    * ``attention_logit[etype](src, dst, edges)`` — edge logits (include
      ``1/sqrt(d)`` here, or set ``scale``).
    * Softmax over neighbors, weighted sum, cross ``sum`` across etypes.
    * ``target_apply[ntype]`` — target-type output projection.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """

    def maybe_scale(logit_fn: LogitFn) -> LogitFn:
        if scale is None:
            return logit_fn
        inv = math.sqrt(float(scale))

        def scaled(src, dst, edges):
            return logit_fn(src, dst, edges) / inv

        return scaled

    etype_dict = {
        etype: RelationSpec(
            message_fn=_src_message(message_apply[etype]),
            reduce="sum",
            attention_logit_fn=maybe_scale(attention_logit[etype]),
            attention_reduce_fn=attention_weight_messages,
        )
        for etype in message_apply
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(target_apply[ntype](mail.nodes[ntype]))
        for ntype in target_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

relational_graph_convolution 

relational_graph_convolution(
    graph,
    relation_apply,
    self_apply,
    *,
    activation=_relu,
    reducer="mean",
)

R-GCN (Relational Graph Convolutional Network) layer.

Schlichtkrull et al., ESWC 2018. For each relation r, messages are relation_apply[r](h_src), neighborhood-aggregated with reducer (mean ≈ 1/|N_r(i)|), then cross-summed. Destinations update as activation(self_apply[n](h) + mailbox).

Parameters:

Name Type Description Default
graph HeteroGraphsTuple

Input heterograph.

 required
relation_apply Mapping[CanonicalEtype, LinearFn]

Per-etype h_src -> message (usually a linear).

 required
self_apply Mapping[Ntype, LinearFn]

Per-ntype self / root term (W_0 in the paper).

 required
activation ActivationFn

Pointwise nonlinearity (default ReLU).

 _relu
reducer str

Per-relation segment reduce (mean or sum).

 'mean'
Source code in anytensor/hetero/models.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@cache
def relational_graph_convolution(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    self_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _relu,
    reducer: str = "mean",
) -> HeteroGraphsTuple:
    """R-GCN (Relational Graph Convolutional Network) layer.

    Schlichtkrull et al., ESWC 2018. For each relation ``r``, messages are
    ``relation_apply[r](h_src)``, neighborhood-aggregated with ``reducer``
    (``mean`` ≈ ``1/|N_r(i)|``), then cross-summed. Destinations update as
    ``activation(self_apply[n](h) + mailbox)``.

    Args:
        graph: Input heterograph.
        relation_apply: Per-etype ``h_src -> message`` (usually a linear).
        self_apply: Per-ntype self / root term (``W_0`` in the paper).
        activation: Pointwise nonlinearity (default ReLU).
        reducer: Per-relation segment reduce (``mean`` or ``sum``).
    """
    etype_dict = {
        etype: RelationSpec(message_fn=_src_message(fn), reduce=reducer)  # type: ignore[arg-type]
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(self_apply[ntype](graph.nodes[ntype]) + mail.nodes[ntype])
        for ntype in self_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

anytensor.hetero.models

Heterogeneous GNN layers as plain functions on :class:HeteroGraphsTuple.

Each function takes a graph plus callables / arrays for the learnable pieces. Weight ownership stays in your framework (Flax, Haiku, torch.nn, NumPy prototypes) — pass lambda x: x @ W or a module __call__ as needed.

Citations (acronym → full name)

  • R-GCN (Relational Graph Convolutional Network) — Schlichtkrull et al., “Modeling Relational Data with Graph Convolutional Networks,” ESWC 2018. https://arxiv.org/abs/1703.06103
  • GraphSAGE (SAmple and aggreGatE; hetero wrap) — Hamilton et al., “Inductive Representation Learning on Large Graphs,” NeurIPS 2017. https://arxiv.org/abs/1706.02216
  • HAN (Heterogeneous Graph Attention Network) — Wang et al., WWW 2019. https://arxiv.org/abs/1903.07293
  • HGT (Heterogeneous Graph Transformer) — Hu et al., WWW 2020. https://arxiv.org/abs/2003.01332
  • CompGCN (Composition-based Multi-Relational GCN) — Vashishth et al., ICLR 2020. https://arxiv.org/abs/1911.03082
  • GAT (Graph Attention Network) edge scores — Veličković et al., ICLR 2018. https://arxiv.org/abs/1710.10903 (used by :func:gat_attention_logit / HAN)

relational_graph_convolution 

relational_graph_convolution(
    graph,
    relation_apply,
    self_apply,
    *,
    activation=_relu,
    reducer="mean",
)

R-GCN (Relational Graph Convolutional Network) layer.

Schlichtkrull et al., ESWC 2018. For each relation r, messages are relation_apply[r](h_src), neighborhood-aggregated with reducer (mean ≈ 1/|N_r(i)|), then cross-summed. Destinations update as activation(self_apply[n](h) + mailbox).

Parameters:

Name Type Description Default
graph HeteroGraphsTuple

Input heterograph.

 required
relation_apply Mapping[CanonicalEtype, LinearFn]

Per-etype h_src -> message (usually a linear).

 required
self_apply Mapping[Ntype, LinearFn]

Per-ntype self / root term (W_0 in the paper).

 required
activation ActivationFn

Pointwise nonlinearity (default ReLU).

 _relu
reducer str

Per-relation segment reduce (mean or sum).

 'mean'
Source code in anytensor/hetero/models.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@cache
def relational_graph_convolution(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    self_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _relu,
    reducer: str = "mean",
) -> HeteroGraphsTuple:
    """R-GCN (Relational Graph Convolutional Network) layer.

    Schlichtkrull et al., ESWC 2018. For each relation ``r``, messages are
    ``relation_apply[r](h_src)``, neighborhood-aggregated with ``reducer``
    (``mean`` ≈ ``1/|N_r(i)|``), then cross-summed. Destinations update as
    ``activation(self_apply[n](h) + mailbox)``.

    Args:
        graph: Input heterograph.
        relation_apply: Per-etype ``h_src -> message`` (usually a linear).
        self_apply: Per-ntype self / root term (``W_0`` in the paper).
        activation: Pointwise nonlinearity (default ReLU).
        reducer: Per-relation segment reduce (``mean`` or ``sum``).
    """
    etype_dict = {
        etype: RelationSpec(message_fn=_src_message(fn), reduce=reducer)  # type: ignore[arg-type]
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(self_apply[ntype](graph.nodes[ntype]) + mail.nodes[ntype])
        for ntype in self_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

hetero_sage 

hetero_sage(
    graph,
    relation_apply,
    combine_apply,
    *,
    activation=_relu,
)

Heterogeneous GraphSAGE mean layer (Hamilton et al., NeurIPS 2017).

GraphSAGE (SAmple and aggreGatE): per-relation map on sources, mean aggregate, cross sum, then activation(combine_apply[n](concat[h_self, mailbox])). Apply is @cache (same pattern as GraphNetwork).

Source code in anytensor/hetero/models.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@cache
def hetero_sage(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    combine_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _relu,
) -> HeteroGraphsTuple:
    """Heterogeneous GraphSAGE mean layer (Hamilton et al., NeurIPS 2017).

    GraphSAGE (SAmple and aggreGatE): per-relation map on sources, ``mean``
    aggregate, cross ``sum``, then
    ``activation(combine_apply[n](concat[h_self, mailbox]))``.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """
    etype_dict = {
        etype: RelationSpec(message_fn=_src_message(fn), reduce="mean")
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(
            combine(concatenate([graph.nodes[ntype], mail.nodes[ntype]], axis=-1))
        )
        for ntype, combine in combine_apply.items()
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

comp_gcn 

comp_gcn(
    graph,
    relation_apply,
    self_apply,
    *,
    composition="mult",
    activation=_relu,
    reducer="sum",
)

CompGCN (Composition-based Multi-Relational GCN) layer.

Vashishth et al., ICLR 2020. Requires edge features on each used etype.

  • mult — relation_apply[r](h_src * e)
  • sum — relation_apply[r](h_src + e)

Then segment reducer, cross sum, and activation(self_apply[n](h) + mailbox).

Source code in anytensor/hetero/models.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
@cache
def comp_gcn(
    graph: HeteroGraphsTuple,
    relation_apply: Mapping[CanonicalEtype, LinearFn],
    self_apply: Mapping[Ntype, LinearFn],
    *,
    composition: str = "mult",
    activation: ActivationFn = _relu,
    reducer: str = "sum",
) -> HeteroGraphsTuple:
    """CompGCN (Composition-based Multi-Relational GCN) layer.

    Vashishth et al., ICLR 2020. Requires edge features on each used etype.

    * ``mult`` — ``relation_apply[r](h_src * e)``
    * ``sum`` — ``relation_apply[r](h_src + e)``

    Then segment ``reducer``, cross ``sum``, and
    ``activation(self_apply[n](h) + mailbox)``.
    """
    if composition not in ("mult", "sum"):
        raise ValueError("composition must be 'mult' or 'sum'")

    def make_msg(apply: LinearFn):
        def msg(src, dst, edges):
            del dst
            if edges is None:
                raise ValueError("comp_gcn requires edge features on every etype")
            composed = src * edges if composition == "mult" else src + edges
            return apply(composed)

        return msg

    etype_dict = {
        etype: RelationSpec(message_fn=make_msg(fn), reduce=reducer)  # type: ignore[arg-type]
        for etype, fn in relation_apply.items()
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(self_apply[ntype](graph.nodes[ntype]) + mail.nodes[ntype])
        for ntype in self_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

han 

han(
    graph,
    meta_path_etypes,
    node_message,
    node_attention_logit,
    semantic_project,
    semantic_query,
    *,
    node_activation=_relu,
    semantic_activation=_tanh,
)

HAN (Heterogeneous Graph Attention Network) layer.

Wang et al., WWW 2019: node-level + semantic attention. Each meta_path_etypes entry is a meta-path hop already stored as a canonical etype (precompute longer paths as their own etypes).

  1. Node-level attention — node_message[e](h_src), logits from node_attention_logit[e](src, dst, edges), softmax over neighbors (GAT-style).
  2. Mailboxes stacked; semantic attention mixes path embeddings with semantic_query after semantic_project. Apply is @cache (same pattern as GraphNetwork).
Source code in anytensor/hetero/models.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
@cache
def han(
    graph: HeteroGraphsTuple,
    meta_path_etypes: Sequence[CanonicalEtype],
    node_message: Mapping[CanonicalEtype, LinearFn],
    node_attention_logit: Mapping[CanonicalEtype, LogitFn],
    semantic_project: LinearFn,
    semantic_query: ArrayTree,
    *,
    node_activation: ActivationFn = _relu,
    semantic_activation: ActivationFn = _tanh,
) -> HeteroGraphsTuple:
    """HAN (Heterogeneous Graph Attention Network) layer.

    Wang et al., WWW 2019: node-level + semantic attention. Each
    ``meta_path_etypes`` entry is a meta-path hop already stored as a
    canonical etype (precompute longer paths as their own etypes).

    1. **Node-level attention** — ``node_message[e](h_src)``, logits from
       ``node_attention_logit[e](src, dst, edges)``, softmax over neighbors
       (GAT-style).
    2. Mailboxes **stacked**; **semantic attention** mixes path embeddings
       with ``semantic_query`` after ``semantic_project``.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """
    if not meta_path_etypes:
        raise ValueError("han requires at least one meta-path etype")

    etype_dict = {
        etype: RelationSpec(
            message_fn=_src_message(node_message[etype]),
            reduce="sum",
            attention_logit_fn=node_attention_logit[etype],
            attention_reduce_fn=attention_weight_messages,
        )
        for etype in meta_path_etypes
    }
    stacked = multi_update_all(graph, etype_dict, cross_reducer="stack")
    nodes = dict(graph.nodes)
    for ntype in {e[2] for e in meta_path_etypes}:
        h_stack = stacked.nodes[ntype]  # (n, R, d)
        h_act = node_activation(h_stack)
        n, r, d = shape(h_act)[:3]
        flat = reshape(h_act, (-1, d))
        proj = semantic_project(flat)
        d_s = shape(proj)[-1]
        proj = reshape(proj, (n, r, d_s))
        proj = semantic_activation(proj)
        q_vec = reshape(semantic_query, (d_s,))
        score = at_sum(proj * q_vec, axes=-1)
        alpha = reshape(_softmax_axis1(score), (n, r, 1))
        nodes[ntype] = at_sum(h_act * alpha, axes=1)
    return graph.update(nodes=nodes)

hgt 

hgt(
    graph,
    message_apply,
    attention_logit,
    target_apply,
    *,
    activation=_identity,
    scale=None,
)

HGT (Heterogeneous Graph Transformer) style layer.

Hu et al., WWW 2020 — typed attention + target projection. Full HGT uses typed query/key/value and edge-type matrices (often multi-head). Fold those into the callables you pass:

  • message_apply[etype](h_src) — value / message projection.
  • attention_logit[etype](src, dst, edges) — edge logits (include 1/sqrt(d) here, or set scale).
  • Softmax over neighbors, weighted sum, cross sum across etypes.
  • target_apply[ntype] — target-type output projection. Apply is @cache (same pattern as GraphNetwork).
Source code in anytensor/hetero/models.py 
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
@cache
def hgt(
    graph: HeteroGraphsTuple,
    message_apply: Mapping[CanonicalEtype, LinearFn],
    attention_logit: Mapping[CanonicalEtype, LogitFn],
    target_apply: Mapping[Ntype, LinearFn],
    *,
    activation: ActivationFn = _identity,
    scale: Optional[float] = None,
) -> HeteroGraphsTuple:
    """HGT (Heterogeneous Graph Transformer) style layer.

    Hu et al., WWW 2020 — typed attention + target projection. Full HGT uses
    typed query/key/value and edge-type matrices (often multi-head). Fold
    those into the callables you pass:

    * ``message_apply[etype](h_src)`` — value / message projection.
    * ``attention_logit[etype](src, dst, edges)`` — edge logits (include
      ``1/sqrt(d)`` here, or set ``scale``).
    * Softmax over neighbors, weighted sum, cross ``sum`` across etypes.
    * ``target_apply[ntype]`` — target-type output projection.
    Apply is ``@cache`` (same pattern as GraphNetwork).
    """

    def maybe_scale(logit_fn: LogitFn) -> LogitFn:
        if scale is None:
            return logit_fn
        inv = math.sqrt(float(scale))

        def scaled(src, dst, edges):
            return logit_fn(src, dst, edges) / inv

        return scaled

    etype_dict = {
        etype: RelationSpec(
            message_fn=_src_message(message_apply[etype]),
            reduce="sum",
            attention_logit_fn=maybe_scale(attention_logit[etype]),
            attention_reduce_fn=attention_weight_messages,
        )
        for etype in message_apply
    }
    mail = multi_update_all(graph, etype_dict, cross_reducer="sum")
    nodes = {
        ntype: activation(target_apply[ntype](mail.nodes[ntype]))
        for ntype in target_apply
    }
    merged = dict(graph.nodes)
    merged.update(nodes)
    return graph.update(nodes=merged)

gat_attention_logit 

gat_attention_logit(
    src, dst, attn_vec_apply, *, negative_slope=0.2
)

GAT-style edge score: LeakyReLU(attn_vec_apply(concat(src, dst))).

Graph Attention Network (GAT; Veličković et al., ICLR 2018) neighborhood scoring. attn_vec_apply maps concatenated features to shape (E, 1) or (E,). Typical use inside a logit callable::

lambda s, d, e: gat_attention_logit(s, d, my_linear)
Source code in anytensor/hetero/models.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def gat_attention_logit(
    src: ArrayTree,
    dst: ArrayTree,
    attn_vec_apply: LinearFn,
    *,
    negative_slope: float = 0.2,
) -> ArrayTree:
    """GAT-style edge score: ``LeakyReLU(attn_vec_apply(concat(src, dst)))``.

    Graph Attention Network (GAT; Veličković et al., ICLR 2018) neighborhood
    scoring. ``attn_vec_apply`` maps concatenated features to shape ``(E, 1)``
    or ``(E,)``. Typical use inside a logit callable::

        lambda s, d, e: gat_attention_logit(s, d, my_linear)
    """
    cat = concatenate([src, dst], axis=-1)
    x = attn_vec_apply(cat)
    return where(x > 0, x, x * negative_slope)