treeflow.traversal.preorder module

treeflow.traversal.preorder.preorder_traversal(topology: TensorflowTreeTopology, mapping: Callable[[TOutputStructure, TInputStructure], TOutputStructure], input: TInputStructure, root_init: TOutputStructure, unroll: str = 'auto') → TOutputStructure

Preorder (parents-before-children) traversal over a fixed topology.

Carries an arbitrary tf.nest output structure per internal node and applies mapping at each non-root internal node. General enough to host the node-height ratio transform (input = per-node (ratio, anchor), root_init = the root height, mapping = the affine update).

Parameters:

  • topology – The tree topology providing preorder_node_indices and parent_indices.

  • mapping – (parent_output, node_input) -> node_output applied at each non-root internal node.

  • input – Per-internal-node input structure, indexed input[node_index - taxon_count].

  • root_init – Output structure for the root node, used to seed the traversal.

  • unroll –

    Which traversal strategy to use – one of "auto" (default), "unrolled", "tensorarray" or "while_loop":

    • "unrolled": a straight-line Python-unrolled graph with no TensorArray. Requires the topology index values to be statically known (tf.get_static_value folds them). Fastest overall.

    • "tensorarray": a Python-unrolled loop that writes a TensorArray. Requires only the node count (index shape) to be static.

    • "while_loop": an AutoGraph tf.while_loop over a TensorArray. Requires nothing static; O(1) graph, for a varying or very large topology.

    • "auto": "unrolled" if values static, else "tensorarray" if count static, else "while_loop".

    XLA / jit_compile notes (empirically verified): "unrolled" is the only mode that XLA-compiles for value+gradient – under jit_compile always use it (a static-value topology), because the TensorArray modes’ backward pass materialises a TensorList that cannot cross the XLA/TF boundary. The TensorArray modes XLA-compile only the forward pass. Wrapping the TensorArray traversal in an inner tf.function does not shield it from an enclosing jit_compile (compilation propagates through).