Dynamic Graph Fields

Use dynamic graph fields when you need per-graph attributes collated during native batching.

This is a standalone program and its captured output for this version:

This page fragment is generated from docs/snippets/dynamic_graph_fields.py. It is intended to show the exact code and typical output for this version.

Commit: 390f8090430944521df7c322996af8aa89f78469

Program

from __future__ import annotations

import mifrost
from mifrost.graph_fields import DType, Mode

from ._helpers import first_successor, load_problem, print_encoding_summary, to_py_list


def main() -> None:
    space, domain, _problem, state0 = load_problem("blocks", "smedium")
    _action, state1 = first_successor(space, state0)

    encoder = mifrost.HGraphEncoder(domain)

    enc0 = encoder.encode(state0)
    enc1 = encoder.encode(state1)
    enc0.goal_distance = 0
    enc0.target_indices = [0, 1, 2]
    enc1.goal_distance = 1
    enc1.target_indices = [0]

    batch_enc = mifrost.batch_encodings(
        [enc0, enc1],
        collate_spec={
            "goal_distance": {"mode": Mode.STACK.value, "dtype": DType.F32.value},
            "target_indices": {
                "mode": Mode.RAGGED_CAT.value,
                "dtype": DType.I64.value,
            },
        },
    )
    print_encoding_summary(
        batch_enc, label="Dynamic attrs via encode + batch_encodings(collate_spec=...)"
    )

    payload = batch_enc.as_dict()
    tensors = payload.get("tensors", {})
    print(
        "native_graph_field_keys_present:",
        sorted(k for k in tensors.keys() if str(k).startswith("__graph__/")),
    )
    print("goal_distance:", to_py_list(batch_enc.goal_distance, max_items=10))
    print("target_indices:", to_py_list(batch_enc.target_indices, max_items=10))
    print("target_indices_ptr:", to_py_list(batch_enc.target_indices_ptr, max_items=10))
    print("collate_spec:", batch_enc.collate_spec())


if __name__ == "__main__":
    main()

Output

== Dynamic attrs via encode + batch_encodings(collate_spec=...) ==
python: 3.12.13
platform: Linux-6.17.0-1018-azure-x86_64-with-glibc2.39
graph_kind: hetero
num_graphs: 2
num_nodes: 24
num_edges: 50
schema_fingerprint: 2869292577121628696
node_types_count: 28
edge_types_count: 54
node_types_head: ['[+]clear[g]', '[+]clear[g][sat]', '[+]handempty[g]', '[+]handempty[g][sat]', '[+]holding[g]', '[+]holding[g][sat]', '[+]on[g]', '[+]on[g][sat]', '[+]ontable[g]', '[+]ontable[g][sat]']
edge_types_head: ['[+]clear[g][sat]|0|_symbol_', '[+]clear[g]|0|_symbol_', '[+]holding[g][sat]|0|_symbol_', '[+]holding[g]|0|_symbol_', '[+]on[g][sat]|0|_symbol_', '[+]on[g][sat]|1|_symbol_', '[+]on[g]|0|_symbol_', '[+]on[g]|1|_symbol_', '[+]ontable[g][sat]|0|_symbol_', '[+]ontable[g]|0|_symbol_']
SNAPSHOT_JSON_BEGIN
{
  "edge_types_count": 54,
  "edge_types_head": [
    "[+]clear[g][sat]|0|_symbol_",
    "[+]clear[g]|0|_symbol_",
    "[+]holding[g][sat]|0|_symbol_",
    "[+]holding[g]|0|_symbol_",
    "[+]on[g][sat]|0|_symbol_",
    "[+]on[g][sat]|1|_symbol_",
    "[+]on[g]|0|_symbol_",
    "[+]on[g]|1|_symbol_",
    "[+]ontable[g][sat]|0|_symbol_",
    "[+]ontable[g]|0|_symbol_"
  ],
  "graph_kind": "hetero",
  "label": "Dynamic attrs via encode + batch_encodings(collate_spec=...)",
  "node_types_count": 28,
  "node_types_head": [
    "[+]clear[g]",
    "[+]clear[g][sat]",
    "[+]handempty[g]",
    "[+]handempty[g][sat]",
    "[+]holding[g]",
    "[+]holding[g][sat]",
    "[+]on[g]",
    "[+]on[g][sat]",
    "[+]ontable[g]",
    "[+]ontable[g][sat]"
  ],
  "num_edges": 50,
  "num_graphs": 2,
  "num_nodes": 24,
  "platform": "Linux-6.17.0-1018-azure-x86_64-with-glibc2.39",
  "python": "3.12.13",
  "schema_fingerprint": 2869292577121628696
}
SNAPSHOT_JSON_END
native_graph_field_keys_present: []
goal_distance: [0.0, 1.0]
target_indices: [0, 1, 2, 0]
target_indices_ptr: [0, 3, 4]
collate_spec: {'goal_distance': {'dtype': 'f32', 'mode': 'stack', 'dim': 1, 'cat_dim': 0, 'inc': {'kind': 'none'}}, 'target_indices': {'dtype': 'i64', 'mode': 'ragged_cat', 'dim': 1, 'cat_dim': 0, 'inc': {'kind': 'none'}}}

Ragged fields expose <attr>_ptr tensors (for example target_indices_ptr) in PyG conversion.

Current Native Behavior

• Native graph fields on BatchEncoding are accessible as attributes (encoding.goal_distance, encoding.target_indices).
• For ragged fields, <key>_ptr is reserved and direct assignment to it is rejected.
• In-place mutation on value tensors returned from native access is write-through.
• Ptr tensors are returned as snapshots (mutating the returned ptr tensor does not mutate native storage).

Collision Rules

Python-side graph-field collation specs (dtype="pyobj") may not reuse native graph-field keys. Attempting to register a colliding key now raises an explicit ValueError.