mattjj

from __future__ import annotations
from contextlib import contextmanager
from typing import NamedTuple, Callable, Optional, Any
import numpy as np

Array = Any

class Node(NamedTuple):
  vjp: Optional[Callable]
  parents: List[Node]