treeflow.cli.benchmark module

treeflow.cli.benchmark.time_fn(f: Callable[[...], T]) → Callable[[int, Iterable[Tensor]], Tuple[float, T]]

treeflow.cli.benchmark.get_tree_likelihood_computation(tree: TensorflowRootedTree, input: Alignment, dtype: DType, bito_instance: object | None = None) → Tuple[Callable[[Tensor], Tensor], Tensor]

treeflow.cli.benchmark.get_tree_likelihood_gtr_computation(tree: TensorflowRootedTree, input: Alignment, dtype: DType, bito_instance: object | None = None)

treeflow.cli.benchmark.get_ratio_transform_computation(tree: TensorflowRootedTree, dtype: DType, bito_instance: object | None = None) → Tuple[Callable[[Tensor], Tensor], Tensor]

treeflow.cli.benchmark.get_ratio_transform_jacobian_computation(tree: TensorflowRootedTree, dtype: DType) → Tuple[Callable[[Tensor], Tensor], Tensor]

treeflow.cli.benchmark.get_constant_coalescent_computation(base_tree: TensorflowRootedTree, dtype: DType) → Tuple[Callable[[Tensor, Tensor], Tensor], Tuple[Tensor, Tensor]]

treeflow.cli.benchmark.get_gradient_fn(task_fn: Callable[[...], Tensor], output_gradients: Tensor | None = None) → Callable[[...], Tensor]

treeflow.cli.benchmark.benchmark(tree: TensorflowRootedTree, input: Alignment, dtype: DType, scaler: Tensor, computation: str, task: str, jit: bool, precompile: bool, replicates: int, bito_instance: object | None = None) → Tuple[float, float | None]