afnio

class Variable:
    """
    A class to represent generic data, such as textual inputs, outputs, or numeric data.

    Attributes:
        data: The raw data, which can be a scalar string or number, or a sequence of
            such scalars.
        role: A specific description of the role of the variable in the agent that
            provides context to the optimizer.
        requires_grad: Whether to track operations for automatic differentiation and
            compute gradients.
        grad: Stores the gradient of the variable, if `requires_grad` is set to `True`
            and backpropagation has been performed.
    """

    # NOTE: Using forward references for `Variable` and `Node` in type hints below to
    #       avoid runtime circular imports

    _data: Optional[Union[str, int, float, List[Union[str, int, float]]]]
    role: str
    """
    A specific description of the role of the variable in the agent. For example, it
    could be "system prompt for sentiment classification", "input email or message",
    "output categories", etc.

    It is used to give more context about the purpose of the Variable to the
    [Optimizer][...optim.optimizer.Optimizer].
    """
    requires_grad: bool
    """
    Is ``True`` if gradients need to be computed for this Variable, ``False`` otherwise.

    Note:
        The fact that gradients need to be computed for a Variable does not mean that
        the [`grad`][..grad] attribute will be populated, see [`is_leaf`][..is_leaf] for
        more details.
    """
    _grad: List["Variable"]

    # TODO: Consider having `VariableMeta` class with `grad_fn` and `output_nr`
    #       as attributes
    _output_nr: Optional[int]
    _grad_fn: Optional["Node"]

    _retain_grad: bool
    is_leaf: bool
    """All Variables that have [`requires_grad`][..requires_grad] which is `False`
    will be leaf Variables by convention.

    For Variables that have [`requires_grad`][..requires_grad] which is ``True``, they
    will be leaf Variables if they were created by the user. This means that they are
    not the result of an operation and so `grad_fn` is None.

    Only leaf Variables will have their [`grad`][..grad] populated during a call to
    [`backward`][..backward]. To get [`grad`][..grad] populated for non-leaf Variables,
    you can use [`retain_grad`][..retain_grad].

    Examples:
        >>> a = afnio.Variable("abc", requires_grad=True)
        >>> a.is_leaf
        True
        >>> b = afnio.Variable("abc", requires_grad=True).upper()
        >>> b.is_leaf
        False
        # b was created by the operation that converts all string characters to uppercase
        >>> c = afnio.Variable("abc", requires_grad=True) + "def"
        >>> c.is_leaf
        False
        # c was created by the addition operation
        >>> d = afnio.Variable("abc").upper()
        >>> d.is_leaf
        True
        # d does not require gradients and so has no operation creating it (that is tracked by the autodiff engine)
        >>> e = afnio.Variable("abc").upper().requires_grad_()
        >>> e.is_leaf
        True
        # e requires gradients and has no operations creating it
    """  # noqa: E501
    variable_id: Optional[str]
    _initialized: bool
    _pending_grad_fn_id: Optional[str]
    _pending_grad: Optional[bool]
    _pending_data: Optional[bool]

    def __init__(
        self,
        data: Optional[Union[str, int, float, List[Union[str, int, float]]]] = "",
        role: str = "",
        requires_grad: bool = False,
    ):
        """
        Initializes a `Variable` instance.

        Args:
            data: The raw data for the variable, which can be a scalar string or number,
                or a sequence of such scalars.
            role: A specific description of the role of the variable in the agent that
                provides context to the [`Optimizer`][...optim.optimizer.Optimizer].
            requires_grad: Whether to track operations for automatic differentiation
                and compute gradients.
        """
        # Websocket attributes
        self.variable_id = None
        self._initialized = False  # Falgs variable is ready to send websocket updates
        self._pending_grad_fn_id = None  # Flags grad_fn is being set (fwd pass running)
        self._pending_grad = False  # Flags grad is being set (bwd pass running)
        self._pending_data = False  # Flags data is being set (optim step running)
        # Internal attributes
        self._data = _validate_and_normalize_data(data)
        self.role = role
        self.requires_grad = requires_grad
        self._retain_grad = False
        self._grad = []
        self._output_nr = 0
        self._grad_fn = None
        self.is_leaf = not requires_grad or self.grad_fn is None

        # Share the variable with the websocket server
        if not is_variable_notify_suppressed():
            try:
                from afnio.cognitive.parameter import Parameter

                # Get the singleton websocket client
                _, ws_client = get_default_clients()

                payload = {
                    "data": self.data,
                    "role": self.role,
                    "requires_grad": self.requires_grad,
                    "obj_type": (
                        "__parameter__"
                        if isinstance(self, Parameter)
                        else "__variable__"
                    ),
                }
                response = run_in_background_loop(
                    ws_client.call("create_variable", payload)
                )
                if "error" in response:
                    raise RuntimeError(
                        response["error"]["data"].get("exception", response["error"])
                    )

                logger.debug(f"Variable created and shared with the server: {self!r}")
                variable_id = response.get("result", {}).get("variable_id")
                if not variable_id:
                    raise RuntimeError(
                        f"Server did not return a variable_id "
                        f"for payload: {payload!r}, response: {response!r}"
                    )
                self.variable_id = variable_id
                self._initialized = True
                register_variable(self)
            except Exception as e:
                logger.error(f"Failed to share Variable with the server: {e}")
                raise

    # TODO: pretty print data lists
    def __repr__(self):
        if self._grad_fn:
            return f"Variable(data={self.data}, role={self.role}, grad_fn={self._grad_fn.name()})"  # noqa: E501
        return f"Variable(data={self.data}, role={self.role}, requires_grad={self.requires_grad})"  # noqa: E501

    # TODO: pretty print data lists
    def __str__(self):

        # Helper function to truncate a string if it's longer than 40 characters
        def truncate_str(s):
            if isinstance(s, (int, float)):
                return str(s)
            if len(s) > 40:
                return f"{s[:20]}...{s[-20:]}"
            return s

        # Helper function to show the first and last three elements if it is long
        def format_list(data_list):
            if len(data_list) > 6:
                truncated = [
                    truncate_str(d) for d in (data_list[:3] + ["..."] + data_list[-3:])
                ]
                return f"[{', '.join(truncated)}]"
            return f"[{', '.join(truncate_str(d) for d in data_list)}]"

        if isinstance(self.data, list):
            data_repr = format_list(self.data)
        else:
            data_repr = truncate_str(self.data)

        if self._grad_fn:
            return f"variable({data_repr}, role={truncate_str(self.role)}, grad_fn={self._grad_fn.name()})"  # noqa: E501
        return f"variable({data_repr}, role={truncate_str(self.role)}, requires_grad={self.requires_grad})"  # noqa: E501

    def __add__(self, other) -> "Variable":
        if not isinstance(other, Variable):
            raise TypeError("Only Variables can be added to each other.")

        from afnio.autodiff.basic_ops import Add

        return Add.apply(self, other)

    def __iadd__(self, other) -> "Variable":
        if not isinstance(other, Variable):
            raise TypeError("Only Variables can be added to each other.")

        from afnio.autodiff.basic_ops import Add

        result = Add.apply(self, other)

        self.data = result.data
        self.role = result.role
        self.requires_grad = result.requires_grad

        # Update the grad function in case `other` also has `requires_grad`
        if result.requires_grad:
            with _allow_grad_fn_assignment():
                self.grad_fn = result.grad_fn

        return self

    def backward(
        self,
        gradient: Optional["Variable"] = None,
        retain_graph: Optional[bool] = None,
        create_graph: Optional[bool] = False,
        inputs: Optional[Union["Variable", Sequence["Variable"]]] = None,
    ) -> None:
        """Computes the gradient of current variable with respect to graph leaves.

        The graph is differentiated using the chain rule. If the variable is non-scalar
        (i.e. its data has more than one element) and requires gradient, the function
        additionally requires specifying a `gradient`. It should be a variable with
        data of matching type and shape, that represents the gradient of the
        differentiated function with respect to `self`.

        This function accumulates gradients in the leaf variables; each call to
        `backward` appends new gradient values to the [`grad`][..grad] list. Clear
        existing gradients before calling it again if accumulation is not desired.

        Note:
            When `inputs` are provided, each input must be a leaf variable. If any
            input is not a leaf, a `RuntimeError` is raised.

        Args:
            gradient: The gradient of the function being differentiated with respect to
                `self`. This argument can be omitted if `self` is a scalar.
            retain_graph: If `False`, the graph used to compute the grads will be freed.
                Setting this to `True` retains the graph, allowing for additional
                backward calls on the same graph, useful for example for multi-task
                learning where you have multiple losses. However, retaining the graph is
                not needed in nearly all cases and can be worked around in a much more
                efficient way. Defaults to the value of `create_graph`.
            create_graph: If `True`, graph of the derivative will be constructed,
                allowing to compute higher order derivative products.
            inputs: Inputs with respect to which the gradient will be accumulated into
                [`grad`][..grad]. All other variables will be ignored.
                If not provided, the gradient is accumulated into all the leaf Variables
                that were used to compute the `Variable`s.

        Raises:
            RuntimeError: If the variable is a leaf (i.e., does not require grad or
                does not have a `grad_fn`).
        """

        if self.is_leaf:
            raise RuntimeError(
                "Variable does not require grad or does not have a grad_fn."
            )

        afnio.autodiff.backward(
            self, gradient, retain_graph, create_graph, inputs=inputs
        )

    def requires_grad_(self, mode: bool = True) -> "Variable":
        """
        Change if autodiff should record operations on this variable: sets this
        variable's [`requires_grad`][..requires_grad] attribute in-place.
        Returns this variable.

        [`requires_grad_`][.]'s main use case is to tell autodiff to begin recording
        operations on a Variable `variable`. If `variable` has `requires_grad=False`
        (because it was obtained through a DataLoader, or required preprocessing or
        initialization), `variable.requires_grad_()` makes it so that autodiff will
        begin to record operations on `variable`.

        Args:
            mode: If autodiff should record operations on this variable.

        Returns:
            self: The variable with updated `requires_grad` and `is_leaf` attributes.

        Examples:
            >>> # Initialize with requires_grad=False for data preprocessing
            >>> x = afnio.Variable(data="abc", role="input")
            >>> x = preprocess(x)  # Preprocess without gradient tracking
            >>> x
            variable(abc, role=input, requires_grad=False)
            ...
            >>> # Now enable requires_grad for backpropagation
            >>> x.requires_grad_()
            >>> output = model(x)
            >>> output.backward()  # Backpropagation through `x`
            >>> x.grad
            variable(ABC, role=input, requires_grad=True)
        """
        self.requires_grad = mode
        self.is_leaf = not self.requires_grad or self.grad_fn is None
        return self

    @property
    def data(self) -> Optional[Union[str, int, float, List[Union[str, int, float]]]]:
        """
        The raw data of the Variable.

        Can be either:

        - a scalar: `str`, `int` or `float`, or
        - a sequence: a `list` or `tuple` of scalars.

        Sequence values must be homogeneous: either all strings or all numbers
        (ints/floats). Tuples are converted to lists; if a numeric sequence mixes
        ints and floats the values are promoted to float. Heterogeneous sequences
        mixing strings and numbers are rejected.

        Returns:
            The underlying data value (a scalar or a list of scalars).
        """  # noqa: E501
        self._wait_for_pending(
            "_pending_data"
        )  # Wait until the pending flag is cleared
        return self._data

    @data.setter
    def data(self, value):
        self._data = _validate_and_normalize_data(value)

    @property
    def output_nr(self) -> int:
        return self._output_nr

    @output_nr.setter
    def output_nr(self, n: int):
        if not isinstance(n, int) or not (n >= 0):
            raise TypeError(
                f"`output_nr` can only be an int greater or equal to 0, "
                f"but {n} is of type {type(n).__name__}"
            )
        self._output_nr = n

    @property
    def grad_fn(self) -> Optional["Node"]:
        self._wait_for_pending(
            "_pending_grad_fn_id"
        )  # Wait until the pending flag is cleared
        return self._grad_fn

    @grad_fn.setter
    def grad_fn(self, fn: Callable):
        """
        Sets the `grad_fn` that will be called by the engine to produce the actual
        gradient for this variable.
        """
        if not getattr(_grad_fn_assignment_allowed, "value", False):
            raise AttributeError(
                "Direct assignment to `grad_fn` is not allowed. "
                "Use Function.apply() to construct Variables with a grad_fn."
            )
        if not self.requires_grad:
            raise RuntimeError(
                "Cannot set `grad_fn` on a variable that does not require gradients. "
                "To enable gradient tracking for this variable, call "
                "`requires_grad_()` before setting `grad_fn`. Only variables with "
                "`requires_grad=True` can have a gradient function (`grad_fn`)."
            )
        self._grad_fn = fn
        self.is_leaf = not self.requires_grad or self.grad_fn is None

    @property
    def grad(self) -> List["Variable"]:
        """
        This attribute is an empty `list` by default and becomes a `list` of
        [`Variable`][..]s the first time a call to [`backward`][..backward] computes
        gradients for `self`. The attribute will then contain the gradients computed and
        future calls to [`backward`][..backward] will accumulate (append) gradients
        into it.

        Raises:
            RuntimeError: If the variable is a non-leaf without `retain_grad` enabled.
        """
        self._wait_for_pending(
            "_pending_grad"
        )  # Wait until the pending flag is cleared
        if self.is_leaf or self._retain_grad:
            return self._grad
        else:
            # Throwing a `UserWarning`` instead of `RuntimeError` could do here, like
            # in Pytorch, but for now I cannot think of any use case for not throwing
            # the error
            raise RuntimeError(
                "Attempted to access `grad` for a non-leaf Variable without "
                "`retain_grad` enabled. Non-leaf Variables do not have their gradients "
                "retained by default in autodiff. To retain gradients for this "
                "Variable, call `retain_grad()` before performing the backward pass."
            )

    @grad.setter
    def grad(self, gradient: List["Variable"]):
        """
        Sets the `grad` for this variable if it is a leaf or has `retain_grad` enabled.
        """
        if not isinstance(gradient, list) or not all(
            isinstance(g, Variable) for g in gradient
        ):
            raise TypeError(
                f"`grad` expects a list of Variables for the gradient to accumulate, "
                f"but got {type(gradient).__name__}."
            )

        if self.is_leaf or self._retain_grad:
            self._grad = gradient
        else:
            # Throwing a `UserWarning`` instead of `RuntimeError` could do here, like
            # in Pytorch, but for now I cannot think of any use case for not throwing
            # the error
            raise RuntimeError(
                "Attempted to set `grad` for a non-leaf Variable without `retain_grad` "
                "enabled. Non-leaf Variables do not have their gradients retained by "
                "default in autodiff. To retain gradients for this Variable, call "
                "`retain_grad()` before performing the backward pass."
            )

    def append_grad(self, gradient: "Variable"):
        """
        Appends a gradient value to the list [`grad`][..grad] for this variable.

        Args:
            gradient: The gradient variable to append.

        Raises:
            RuntimeError: If the variable is a non-leaf without `retain_grad` enabled.
        """
        if self.is_leaf or self._retain_grad:
            self._on_append_grad(gradient)
            self._grad.append(gradient)
        else:
            # Throwing a `UserWarning`` instead of `RuntimeError` could do here, like
            # in Pytorch, but for now I cannot think of any use case for not throwing
            # the error
            raise RuntimeError(
                "Attempted to append to `grad` for a non-leaf Variable without "
                "`retain_grad` enabled. Non-leaf Variables do not have their gradients "
                "retained by default in autodiff. To retain gradients for this "
                "Variable, call `retain_grad()` before performing the backward pass."
            )

    def retain_grad(self):
        """
        Enables this Variable to have their [`grad`][..grad] populated during
        [`backward()`][..backward]. This is a no-op for leaf Variables.

        Raises:
            RuntimeError: If the variable is a leaf.
        """
        if not self.is_leaf:
            self._retain_grad = True
        else:
            raise RuntimeError("Cannot call `retain_grad` on a leaf variable")

    def detach(self) -> "Variable":
        """
        Returns a new Variable, detached from the computation graph.
        This new Variable will not have a `grad_fn` and will not track gradients.

        Returns:
            A new Variable with the same data and role, but with `requires_grad` set to `False`.
        """  # noqa: E501
        return Variable(self.data, role=self.role, requires_grad=False)

    # def clone(self):
    #     """
    #     Create a copy of this Variable, preserving the data.
    #     """
    #     return copy.deepcopy(self)

    def __deepcopy__(self, memo):
        if not self.is_leaf:
            raise RuntimeError(
                "Only Variables created explicitly by the user "
                "(graph leaves) support the deepcopy protocol at the moment."
            )
        if id(self) in memo:
            return memo[id(self)]

        with afnio.no_grad():
            new_variable = Variable(
                data=deepcopy(self.data, memo),
                role=self.role,
                requires_grad=self.requires_grad,
            )

            new_variable._retain_grad = self._retain_grad
            new_variable._output_nr = self._output_nr

            if self.grad_fn:
                with _allow_grad_fn_assignment():
                    new_variable.grad_fn = deepcopy(
                        self.grad_fn, memo
                    )  # Also sets `is_leaf`
            if self.grad != []:
                new_variable.grad = deepcopy(self.grad, memo)

            new_variable.__dict__ = deepcopy(self.__dict__, memo)

            memo[id(self)] = new_variable

            return new_variable

    def copy_(self, src: "Variable") -> "Variable":
        """
        Copies the data from the source Variable into this Variable.

        Args:
            src: The source Variable to copy from.

        Returns:
            self: The current Variable with updated `data`, `role` and `requires_grad`.

        Raises:
            TypeError: If the source is not a Variable.
            ValueError: If the source data type does not match the target data type.
        """
        if not is_variable(src):
            raise TypeError(
                f"Expected `src` to be a Variable, but got {type(src).__name__}."
            )

        is_scalar_self = is_scalar_variable(self)
        is_scalar_src = is_scalar_variable(src)

        if is_scalar_self and is_scalar_src:
            self.data = src.data
        elif not is_scalar_self and not is_scalar_src:
            if len(self.data) != len(src.data):
                raise ValueError(
                    f"Cannot copy list `data` fields of different lengths: "
                    f"{len(self.data)} vs {len(src.data)}."
                )
            self.data = src.data.copy()
        else:
            raise ValueError(
                f"Cannot copy data from {type(src.data).__name__} "
                f"to {type(self.data).__name__}."
            )

        self.role = src.role
        self.requires_grad = src.requires_grad
        return self

    def is_floating_point(self) -> bool:
        """
        Checks if the Variable's [`data`][..data] contains floating-point values.

        Returns:
            `True` if the data is a floating-point type (either scalar or all elements in a list/tuple are floating-point).
        """  # noqa: E501
        if isinstance(self.data, float):
            return True

        if isinstance(self.data, (list, tuple)) and self.data:
            return all(isinstance(d, float) for d in self.data)

        return False

    def to(self, dtype: Optional[type] = None) -> "Variable":
        """
        Cast the [`data`][..data] of the Variable to the specified dtype.

        Args:
            dtype: The target type to cast the data (e.g., `float`, `int`, `str`).

        Returns:
            A new Variable with [`data`][..data] cast to the target `dtype`.
        """
        if dtype is not None:
            if not is_scalar_variable(self):
                # Cast each element in the list to the target dtype
                new_data = [dtype(d) for d in self.data]
            else:
                # Cast scalar data to the target dtype
                new_data = dtype(self.data)
        else:
            # No dtype casting
            new_data = self.data

        # Return a new Variable with the same role and requires_grad, but updated data
        return Variable(data=new_data, role=self.role, requires_grad=self.requires_grad)

    def _on_variable_change(self, field: str, value: Any):
        """
        Notify the server of a change in the variable's attributes.
        This method is called whenever an attribute of the variable is set.
        It sends a notification to the server with the updated field and value.

        Args:
            field: The name of the field that changed.
            value: The new value of the field.

        Raises:
            RuntimeError: If the variable is not registered with the server or if the
                server response does not match the request.
            TypeError: If the provided value is of an unexpected type for the field.
        """
        from afnio._utils import _serialize_arg

        if is_variable_notify_suppressed():
            return  # Do not notify server

        if self.variable_id is None:
            logger.error(
                f"Cannot notify server: "
                f"variable_id=None, field='{field}', value={value!r}"
            )
            raise RuntimeError("Cannot notify server: variable_id is None.")

        if field in {
            "output_nr",
            "grad_fn",
            "grad",
            "_initialized",
            "_pending_grad_fn_id",
            "_pending_grad",
            "_pending_data",
            "__dict__",  # Avoids server error when calling `Optimizer.load_state_dict`
        }:
            # Do not notify for the property setter, as we already notify
            # for all the changes made inside the property setter.
            # Also do not notify for `_initialized` and pending states
            return
        elif field == "_data":
            field = "data"  # `data` is a property only on the client
            end_value = value
        elif field == "_grad":
            if not isinstance(value, list):
                raise TypeError(
                    f"Expected `value` to be a list for field '{field}', "
                    f"but got {type(value).__name__}."
                )
            end_value = [_serialize_arg(g) for g in value]
        elif field == "_grad_fn":
            # Only allow notification if inside the `__iadd__` method
            if not _called_directly_from_iadd():
                raise RuntimeError(
                    "Setting `grad_fn` is only allowed on the server by the autodiff "
                    "engine. Do not use `_allow_grad_fn_assignment()` on the client."
                )
            end_value = value.node_id  # Use only the node ID for notification
        else:
            end_value = value

        payload = {
            "variable_id": self.variable_id,
            "field": field,
            "value": end_value,
        }

        try:
            _, ws_client = get_default_clients()
            response = run_in_background_loop(
                ws_client.call("update_variable", payload)
            )
            if "error" in response:
                raise RuntimeError(
                    response["error"]["data"].get("exception", response["error"])
                )

            # Check server response
            if (
                response["result"]["variable_id"] != self.variable_id
                or response["result"]["field"] != field
                or response["result"]["value"] != end_value
            ):
                raise RuntimeError(
                    f"Server response mismatch: (received {response['result']!r}, "
                    f"but expected variable_id={self.variable_id!r}, field={field!r}, "
                    f"value={end_value!r})"
                )
            logger.debug(
                f"Variable change notified to server and confirmed: "
                f"variable_id={self.variable_id!r}, field='{field}', "
                f"value={end_value!r}"
            )

        except Exception as e:
            logger.exception(f"Failed to notify server of variable change: {e}")
            raise

    def _on_append_grad(self, gradient: "Variable"):
        """
        Notify the server that a new gradient has been appended to this variable.

        This method is called before the gradient is added to the local [`grad`][..grad]
        list. It sends an 'append_grad' RPC request to the server, including the
        variable's ID and the serialized gradient. The method blocks until the server
        acknowledges the append operation, ensuring synchronization between client and
        server.

        Args:
            gradient: The gradient variable to append.

        Raises:
            RuntimeError: If the variable is not registered with the server or if the
                server response does not match the request.
            TypeError: If the provided gradient is not a Variable.
        """
        from afnio._utils import _serialize_arg

        if is_variable_notify_suppressed():
            return  # Do not notify server

        if self.variable_id is None:
            logger.error(
                f"Cannot notify server: variable_id=None, gradient={gradient!r}"
            )
            raise RuntimeError("Cannot notify server: variable_id is None.")

        if not isinstance(gradient, Variable):
            raise TypeError(
                f"Expected `value` to be a Variable, but got {type(gradient).__name__}."
            )

        ser_grad = _serialize_arg(gradient)

        payload = {
            "variable_id": self.variable_id,
            "gradient": ser_grad,
        }

        try:
            _, ws_client = get_default_clients()
            response = run_in_background_loop(ws_client.call("append_grad", payload))
            if "error" in response:
                raise RuntimeError(
                    response["error"]["data"].get("exception", response["error"])
                )

            # Check server response
            if (
                response["result"]["variable_id"] != self.variable_id
                or response["result"]["gradient_id"] != gradient.variable_id
            ):
                raise RuntimeError(
                    f"Server response mismatch: (received {response['result']!r}, "
                    f"but expected variable_id={self.variable_id!r}, "
                    f"gradient={ser_grad!r}"
                )
            logger.debug(
                f"Gradient append notified to server and confirmed: "
                f"variable_id={self.variable_id!r}, gradient={ser_grad!r}"
            )

        except Exception as e:
            logger.exception(f"Failed to notify server of gradient append: {e}")
            raise

    def __setattr__(self, name, value):
        super().__setattr__(name, value)
        if getattr(self, "_initialized", False):
            self._on_variable_change(name, value)
        # TODO: Should we handle the else condition and throw an error?

    def _wait_for_pending(
        self, attr_name: str, timeout: float = 3, interval: float = 0.01
    ) -> None:
        """
        Wait until the attribute specified by `attr_name` is no longer truthy.
        Uses `time.monotonic()` for more reliable timeout measurement.

        Args:
            attr_name: Name of the attribute to wait on.
            timeout: Maximum time to wait, in seconds.
            interval: How frequently to check the attribute, in seconds.

        Raises:
            RuntimeError: If the attribute remains truthy after the timeout.
        """
        end_time = time.monotonic() + timeout
        while getattr(self, attr_name):
            if time.monotonic() > end_time:
                raise RuntimeError(
                    f"Timeout waiting for {attr_name} to be cleared "
                    f"for variable_id={self.variable_id}"
                )
            time.sleep(interval)

class GradientEdge(NamedTuple):
    """
    Object representing a given gradient edge within the autodiff backward graph.

    A GradientEdge identifies *where* a gradient should be propagated during
    backpropagation. It points to a specific backward node and a specific
    output index of that node.

    Each [`Variable`][..Variable] that participates in autodiff is associated with
    one or more GradientEdge objects, which define how gradients flow through the
    computation graph.

    To get the gradient edge where a given Variable gradient will be computed,
    you can do `edge = autodiff.graph.get_gradient_edge(variable)`.

    Attributes:
        node: The backward node responsible for producing gradients.
            This is an instance of a backward Function (e.g. AccumulateGrad,
            SumBackward0, ChatCompletionBackward0).
        output_nr: The index of the output of `node` to which this edge corresponds.
            Required because a backward node may produce multiple gradient
            outputs.
    """

    node: Node
    output_nr: int

    def __repr__(self):
        name = (
            f"<{self.node.name()} object at {hex(id(self.node))}>"
            if self.node
            else "None"
        )
        return f"({name}, {self.output_nr})"

    def __str__(self):
        name = (
            f"<{self.node.name()} object at {hex(id(self.node))}>"
            if self.node
            else "None"
        )
        return f"({name}, {self.output_nr})"

class Node:
    """
    Base class for nodes in the autograd computation graph.

    A Node represents a single operation in the backward graph and is
    responsible for producing gradients during backpropagation. Each Node
    may have multiple inputs and outputs and is connected to other nodes
    via [`GradientEdge`][..GradientEdge] objects.

    Attributes:
        next_functions: The gradient edges to this node's parent nodes
            in the backward graph.
    """

    def __init__(self, next_functions: Optional[Tuple["GradientEdge"]] = None):
        self._next_functions = next_functions if next_functions else ()
        self._name = None
        self.node_id = None

    def __repr__(self):
        return f"<afnio.autodiff.function.{self.name()} object at {hex(id(self))}>"

    def __str__(self):
        return f"<{self.name()} object at {hex(id(self))}>"

    def apply(self, *args):
        raise NotImplementedError("Subclasses should implement this method.")

    def name(self) -> str:
        """
        Returns the name of the backward node.

        Returns:
            The name of the backward node, which is typically the name of the forward operation that produced the output associated with this node.

        Examples:
            >>> import afnio
            >>> import afnio.cognitive.functional as F
            >>> a = afnio.Variable("Hello,", requires_grad=True)
            >>> b = afnio.Variable("world!", requires_grad=True)
            >>> c = F.sum([a, b])
            >>> assert isinstance(c.grad_fn, afnio.autodiff.graph.Node)
            >>> print(c.grad_fn.name())
            SumBackward0
        """  # noqa: E501
        return self._name

    @property
    def next_functions(self) -> Tuple["GradientEdge"]:
        """
        Returns the gradient edges to this node's parent nodes in the backward graph.

        Each entry is a tuple of `(Node, output_nr)` identifying where gradients
        should be propagated next during backpropagation.

        Returns:
            A tuple of [`GradientEdge`][..GradientEdge] objects representing the edges to this node's parent nodes in the backward graph.
        """  # noqa: E501
        return self._next_functions

    @next_functions.setter
    def next_functions(self, edges: Tuple["GradientEdge", ...]):
        self._next_functions = edges