treeflow.model.io module

treeflow.model.io.flatten_tensor_to_1d_slices(name: str, x: Tensor) → Dict[str, Tensor]

Returns

treeflow.model.io.flatten_samples_to_dict(samples: object, distribution: JointDistribution | None = None) → Tuple[Dict[str, Tensor], Dict[str, List[str]]]

treeflow.model.io.calculate_tree_stats(name: str, tree: TensorflowRootedTree) → Dict[str, Tensor]

treeflow.model.io.write_samples_to_file(samples: object, distribution: JointDistribution, fname: str, sep=',', vars: Iterable[str] | None = None, tree_vars: Mapping[str, TensorflowRootedTree] | None = None)