>>> graph.partial_trace(np.diag([1, 0, 0, 0]), [2, 2], 1, name="partial")
<Tensor: name="partial", operation_name="partial_trace", shape=(2, 2)>
>>> result = qctrl.functions.calculate_graph(graph=graph, output_node_names=["partial"])
>>> result.output["partial"]["value"]
array([[[1, 0], [0, 0]])