From 5eae012d0007c02006752a67b8ad8f457ab85149 Mon Sep 17 00:00:00 2001
From: Nicolas Kruse <nicolas@nonan.net>
Date: Tue, 31 Mar 2026 11:34:46 +0200
Subject: [PATCH] Updated quaternion handling for usage with grad() function

---
 src/copapy/_autograd.py   | 10 ++++++++--
 src/copapy/_quaternion.py |  5 ++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/copapy/_autograd.py b/src/copapy/_autograd.py
index 4840805..2413da2 100644
--- a/src/copapy/_autograd.py
+++ b/src/copapy/_autograd.py
@@ -1,3 +1,5 @@
+from copapy._quaternion import quaternion
+
 from . import value, vector, tensor
 import copapy.backend as cpb
 from typing import Any, Sequence, overload
@@ -12,8 +14,10 @@ def grad(x: Any, y: vector[Any]) -> vector[float]: ...
 @overload
 def grad(x: Any, y: tensor[Any]) -> tensor[float]: ...
 @overload
+def grad(x: Any, y: quaternion) -> quaternion: ...
+@overload
 def grad(x: Any, y: Sequence[value[Any]]) -> list[unifloat]: ...
-def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any]) -> Any:
+def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any] | quaternion) -> Any:
     """Returns the partial derivative dx/dy where x needs to be a scalar
     and y might be a scalar, a list of scalars, a vector or matrix. It
     uses automatic differentiation in reverse-mode.
@@ -32,7 +36,7 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any
     if isinstance(y, tensor):
         y_set = {v.get_scalar(0) for v in y.flatten()}
     else:
-        assert isinstance(y, Sequence) or isinstance(y, vector)
+        assert isinstance(y, Sequence) or isinstance(y, vector) or isinstance(y, quaternion)
         y_set = set(y)
 
     edges = cpb.get_all_dag_edges_between([x.net.source], (v.net.source for v in y_set if isinstance(v, value)))
@@ -125,6 +129,8 @@ def grad(x: Any, y: value[Any] | Sequence[value[Any]] | vector[Any] | tensor[Any
         return grad_dict[y.net]
     if isinstance(y, vector):
         return vector(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values)
+    if isinstance(y, quaternion):
+        return quaternion(grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values)
     if isinstance(y, tensor):
         return tensor([grad_dict[yi.net] if isinstance(yi, value) else 0.0 for yi in y.values], y.shape)
     return [grad_dict[yi.net] for yi in y]
diff --git a/src/copapy/_quaternion.py b/src/copapy/_quaternion.py
index 6750a70..b1445dc 100644
--- a/src/copapy/_quaternion.py
+++ b/src/copapy/_quaternion.py
@@ -1,4 +1,4 @@
-from typing import overload, Iterable, Callable, Any
+from typing import overload, Iterable, Callable, Any, Iterator
 from ._vectors import vector
 from ._tensors import tensor
 import copapy as cp
@@ -208,6 +208,9 @@ class quaternion(ArrayType[float]):
         """
         return quaternion(func(x) for x in self.values)
 
+    def __iter__(self) -> Iterator[value[float] | float]:
+        return iter(self.values)
+
     def __neg__(self) -> 'quaternion':
         return quaternion(-self.w, -self.x, -self.y, -self.z)