treeflow.traversal.postorder module

class treeflow.traversal.postorder.PostorderTopologyData(child_indices: Tensor)

Bases: object

child_indices: Tensor

treeflow.traversal.postorder.static_taxon_count(topology) → int | None

The taxon count as a Python int if statically known, else None.

topology.taxon_count is prefer_static (derived from parent_indices’s shape, which tf.function preserves), so it resolves to a concrete value when the node count is known – unlike derived index arrays such as preorder_node_indices, whose static shape a boolean_mask can drop even when the input shape is known.

treeflow.traversal.postorder.postorder_node_traversal(topology: TensorflowTreeTopology, mapping: Callable[[TOutputStructure, TInputStructure, PostorderTopologyData], TOutputStructure], input: TInputStructure, leaf_init: TOutputStructure, xla_compatible: bool = False, unroll: str = 'auto') → TOutputStructure

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

Carries an arbitrary tf.nest output structure per node and applies mapping at each internal node. General enough to host the phylogenetic likelihood (input = per-node child transition matrices, leaf_init = leaf partials, mapping = the Felsenstein combine step; the rescaled variant can carry (partials, log_scale) as the output structure and sum the log_scale component afterwards).

Parameters:

• topology – The tree topology providing postorder_node_indices and child_indices.

• mapping – (child_output, node_input, topology_data) -> node_output applied at each internal node, where child_output stacks the children’s outputs on axis 0.

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

• leaf_init – Per-leaf output structure (leaves on axis 0), used to seed the traversal.

• xla_compatible – Only affects the TensorArray paths ("tensorarray"/"while_loop"). When True the TensorArray is seeded with unstack and children are read with per-child read + stack; when False (default) it is seeded with scatter and children are read with gather. scatter/gather have no XLA kernel, so xla_compatible=True is required to XLA-compile the TensorArray forward pass (see the unroll notes on XLA). It can scale quadratically under XLA, so it is for forward-only XLA use, not a default.

• 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 – true for a constant/eager topology). Fastest overall (linear, no per-step loop overhead; one reverse-sweep gradient).

  • "tensorarray": a Python-unrolled loop that writes a TensorArray. Requires only the node count (index shape) to be static; the index values may be runtime tensors (e.g. a topology routed through a JointDistribution).

  • "while_loop": an AutoGraph tf.while_loop over a TensorArray. Requires nothing static; O(1) graph, so it is the right choice when the topology (or its size) varies per call or is very large.

  • "auto": pick the fastest feasible – "unrolled" if the values are static, else "tensorarray" if the count is 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 "unrolled" (so a static-value topology), since the TensorArray modes’ backward pass materialises a TensorList that cannot cross the XLA/TF boundary.

  • The "tensorarray"/"while_loop" modes XLA-compile only the forward pass, and only with xla_compatible=True (the default scatter/gather ops have no XLA kernel). Their value+gradient never XLA-compiles regardless of xla_compatible.

  • Wrapping a TensorArray traversal in an inner tf.function does not shield it from an enclosing jit_compile: the compilation propagates through (even with jit_compile=False on the inner function), so it is not an escape hatch for the above.