From 9bc8affd00589260567652cf34d0df55c64c3559 Mon Sep 17 00:00:00 2001 From: Mattia Ragni Date: Fri, 19 Jun 2026 13:48:58 +0200 Subject: [PATCH 1/2] Feature: optional forward diff mode --- src/tensorial/gcnn/_diff.py | 36 +++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/src/tensorial/gcnn/_diff.py b/src/tensorial/gcnn/_diff.py index 56311ee..a5dc94c 100644 --- a/src/tensorial/gcnn/_diff.py +++ b/src/tensorial/gcnn/_diff.py @@ -20,7 +20,6 @@ DERIV_DELIMITER: Final[str] = "," ArgumentSpecifier = Union[int, "gcnn.typing.TreePath"] - class DerivableGraphFunction(Protocol): def __call__( self, graph: jraph.GraphsTuple, *args: jt.PyTree @@ -129,14 +128,15 @@ def evaluator( argnum: int = 0, scale: float = 1.0, at: dict | None = None, + mode: str = "rev", ) -> "Evaluator": if at is not None: at = flax.core.FrozenDict(at) - return Evaluator(func, self, return_graph, argnum, scale=scale, at=at) + return Evaluator(func, self, return_graph, argnum, scale=scale, at=at, mode=mode) @abc.abstractmethod def build_derivative_fn( - self, func: DerivableGraphFunction, return_graph: bool, argnum: int + self, func: DerivableGraphFunction, return_graph: bool, argnum: int, mode: str = "rev" ) -> DerivableGraphFunction: """Get evaluate function from derivative""" @@ -277,7 +277,7 @@ def __truediv__(self, other: "GraphEntrySpecLike | SingleDerivative") -> "MultiD return MultiDerivative((self, SingleDerivative.create(self.out, wrt))) def build_derivative_fn( - self, func: DerivableGraphFunction, return_graph: bool, argnum: int + self, func: DerivableGraphFunction, return_graph: bool, argnum: int, mode: str = "rev" ) -> DerivableGraphFunction: if not argnum >= 0: raise ValueError(f"argnum must be >= 0, got: {argnum}") @@ -286,16 +286,21 @@ def build_derivative_fn( # Scalar valued diff_fn = jax.grad else: - # Vector valued - # note: we could use forward mode in cases where dim inputs << outputs - # but the memory use is often exterme - diff_fn = jax.jacrev + # Vector valued; forward mode is more memory-efficient when dim(input) << dim(output) + diff_fn = jax.jacfwd if mode == "fwd" else jax.jacrev def _diff_and_pre_process( graph: jraph.GraphsTuple, *args: jt.PyTree ) -> tuple[jt.Array, jraph.GraphsTuple]: value, graph = func(graph, *args) value, graph = self._pre_process(value, graph) + if mode == "fwd": # TODO: workaround for e3nn IrrepsArray + jacfwd. Remove once upstream handles the appended tangent axis. + # jacfwd appends the tangent axis to the output via out_axes=-1. + # When `value` is an e3nn IrrepsArray, this breaks its shape invariant + # (last dim must match the irreps dimension), raising in the IrrepsArray + # rewrap inside jacfwd's tree_unflatten. Differentiating the raw array + # avoids the issue; axis bookkeeping is handled downstream in _post_process. + value = base.as_array(value) return value, graph do_diff = diff_fn(_diff_and_pre_process, argnums=1 + argnum, has_aux=True) @@ -436,7 +441,7 @@ def __iter__(self): yield from self.parts.__iter__() def build_derivative_fn( - self, func: DerivableGraphFunction, return_graph: bool, argnum: int + self, func: DerivableGraphFunction, return_graph: bool, argnum: int, mode: str = "rev" ) -> DerivableGraphFunction: # Work our way from right to left creating the derivative evaluators argnums = [] @@ -447,9 +452,9 @@ def build_derivative_fn( else: argnums.append(self.graph_tuple_paths[wrt_path]) - func = self[0].build_derivative_fn(func, return_graph=return_graph, argnum=argnums[0]) + func = self[0].build_derivative_fn(func, return_graph=return_graph, argnum=argnums[0], mode=mode) for part, argnum_ in zip(self[1:], argnums[1:]): - func = part.build_derivative_fn(func, return_graph=return_graph, argnum=argnum_) + func = part.build_derivative_fn(func, return_graph=return_graph, argnum=argnum_, mode=mode) return func @@ -478,6 +483,7 @@ class Evaluator: argnum: int scale: float = 1.0 at: flax.core.FrozenDict[str, jt.PyTree] | None = dataclasses.field(default=None, hash=False) + mode: str = "rev" # will be set in __post_init__ _evaluate_at: DerivableGraphFunction = dataclasses.field(init=False) @@ -486,7 +492,7 @@ def __post_init__(self): object.__setattr__( self, "_evaluate_at", - self.spec.build_derivative_fn(self.func, return_graph=True, argnum=self.argnum), + self.spec.build_derivative_fn(self.func, return_graph=True, argnum=self.argnum, mode=self.mode), ) def __call__( @@ -526,6 +532,7 @@ def diff( scale: float = 1.0, at: dict | None = None, return_graph=False, + mode: str = "rev", ) -> Evaluator: """ Constructs a JAX-compatible evaluator for computing single or multiple derivatives @@ -568,6 +575,9 @@ def diff( If True, the derivative tensor is packaged into a new Graph object under the name specified by the 'out' argument. If False, the function returns the raw derivative tensor. Defaults to False. + mode (str): + "rev" (default) uses reverse-mode AD (jax.jacrev). "fwd" uses forward-mode AD + (jax.jacfwd), which is more memory-efficient when dim(input) << dim(output). Returns: Evaluator: A callable object that takes a Graph and returns the computed @@ -599,7 +609,7 @@ def diff( else: deriv = MultiDerivative.create(of, wrt, out) - return deriv.evaluator(deriv.adapt(func), return_graph, scale=scale, at=at) + return deriv.evaluator(deriv.adapt(func), return_graph, scale=scale, at=at, mode=mode) def ordered_unique_indices(lst): From 6fcedb79781fa492c084d3f726ec0b38be38b5b3 Mon Sep 17 00:00:00 2001 From: Mattia Ragni Date: Mon, 22 Jun 2026 17:09:10 +0200 Subject: [PATCH 2/2] Fix: deleted value = base.as_array(value) if fwd mode in _diff.py because the e3nn-jax bug has been fixed in version 0.21.1. Modified accordingly pyproject.toml. --- pyproject.toml | 2 +- src/tensorial/gcnn/_diff.py | 8 +------- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index deefe13..189b79c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ 'jraph', 'orbax-checkpoint', 'pytray', - 'e3nn-jax', + 'e3nn-jax>=0.21.1', "equinox>=0.12.0", "reax>=0.6.10,<0.7", "rich", diff --git a/src/tensorial/gcnn/_diff.py b/src/tensorial/gcnn/_diff.py index a5dc94c..4691aed 100644 --- a/src/tensorial/gcnn/_diff.py +++ b/src/tensorial/gcnn/_diff.py @@ -294,13 +294,7 @@ def _diff_and_pre_process( ) -> tuple[jt.Array, jraph.GraphsTuple]: value, graph = func(graph, *args) value, graph = self._pre_process(value, graph) - if mode == "fwd": # TODO: workaround for e3nn IrrepsArray + jacfwd. Remove once upstream handles the appended tangent axis. - # jacfwd appends the tangent axis to the output via out_axes=-1. - # When `value` is an e3nn IrrepsArray, this breaks its shape invariant - # (last dim must match the irreps dimension), raising in the IrrepsArray - # rewrap inside jacfwd's tree_unflatten. Differentiating the raw array - # avoids the issue; axis bookkeeping is handled downstream in _post_process. - value = base.as_array(value) + return value, graph do_diff = diff_fn(_diff_and_pre_process, argnums=1 + argnum, has_aux=True)