afnio.autodiff.grad_mode

def is_grad_enabled() -> bool:
    """
    Check whether grad mode is currently enabled.

    Returns:
        `True` if grad mode is currently enabled, `False` otherwise.
    """
    return getattr(_grad_enabled, "enabled", True)

def set_grad_enabled(mode: bool):
    """
    Set the global state of gradient calculation on or off.

    `set_grad_enabled` will enable or disable gradients based on its argument `mode`.

    Args:
        mode: If `True`, enables gradient calculation. If `False`, disables it.

    Examples:
        >>> x = afnio.Variable("Hello", requires_grad=True)
        >>> _ = afnio.set_grad_enabled(True)
        >>> y = x + x
        >>> y.requires_grad
        True
        >>> _ = afnio.set_grad_enabled(False)
        >>> y = x + x
        >>> y.requires_grad
        False
    """
    _grad_enabled.enabled = mode

@contextmanager
def no_grad():
    """
    Context manager that disables gradient calculation. All operations within this block
    will not track gradients, making them more memory-efficient.

    Disabling gradient calculation is useful for inference, when you are sure
    that you will not call [`Variable.backward()`][afnio.Variable.backward]. It will
    reduce memory consumption for computations that would otherwise have
    `requires_grad=True`.

    In this mode, the result of every computation will have
    `requires_grad=False`, even when the inputs have `requires_grad=True`.
    There is an exception! All factory functions, or functions that create
    a new Variable and take a requires_grad kwarg, will NOT be affected by
    this mode.

    This context manager is thread local; it will not affect computation
    in other threads.

    Also functions as a decorator.

    Examples:
        >>> x = afnio.Variable("abc", role="variable", requires_grad=True)
        >>> with afnio.no_grad():
        ...     y = x + x
        >>> y.requires_grad
        False
        >>> @afnio.no_grad()
        ... def doubler(x):
        ...     return x + x
        >>> z = doubler(x)
        >>> z.requires_grad
        False
        >>> @afnio.no_grad()
        ... def tripler(x):
        ...     return x + x + x
        >>> z = tripler(x)
        >>> z.requires_grad
        False
        >>> # factory function exception
        >>> with afnio.no_grad():
        ...     a = afnio.cognitive.Parameter("xyz")
        >>> a.requires_grad
        True
    """
    previous_state = is_grad_enabled()  # Store the current state
    set_grad_enabled(False)  # Disable gradients
    try:
        yield  # Execute the block
    finally:
        set_grad_enabled(previous_state)  # Restore the original state