eae-runtime 0.1.0__py3-none-any.whl

eae_runtime/passes/logging_pass.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from ..adjoint import AdjointState
+from ..events import EventBus, EventType
+from .base import EAEPass
+
+
+class LoggingPass(EAEPass):
+    """Emits an ADJOINT_MODIFIED event carrying adjoint.statistics(); does
+    not otherwise touch the adjoint. Demonstrates that observation-only
+    passes are just as first-class as mutating ones."""
+
+    name = "LoggingPass"
+
+    def __init__(self, event_bus: Optional[EventBus] = None):
+        self.event_bus = event_bus or EventBus()
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        stats = adjoint.statistics()
+        self.event_bus.emit(
+            EventType.ADJOINT_MODIFIED,
+            layer_id=adjoint.layer_id,
+            block=adjoint.block,
+            pass_name=self.name,
+            **stats,
+        )
+        return adjoint.clone()

eae_runtime/passes/quantize.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import torch
+
+from ..adjoint import AdjointState
+from .base import EAEPass
+
+
+class FP16Pass(EAEPass):
+    """Downcasts the adjoint to fp16 and back, simulating the numerical
+    effect of a reduced-precision communication/compute path."""
+
+    name = "FP16Pass"
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        new = adjoint.clone()
+        orig_dtype = new.tensor.dtype
+        new.tensor = new.tensor.to(torch.float16).to(orig_dtype)
+        return new
+
+
+class Int8QuantizationPass(EAEPass):
+    """Fake-quantizes the adjoint to int8 with a per-tensor symmetric scale,
+    then immediately dequantizes. This is the standard 'fake quant' trick
+    used to study the effect of low-precision gradient compression without
+    needing real int8 kernels. A production FP8Pass would follow the same
+    shape but call into real fp8 kernels via the Backend Manager."""
+
+    name = "Int8QuantizationPass"
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        new = adjoint.clone()
+        t = new.tensor.float()
+        max_abs = t.abs().max()
+        if max_abs == 0:
+            return new
+        scale = max_abs / 127.0
+        q = torch.clamp(torch.round(t / scale), -127, 127)
+        dq = (q * scale).to(adjoint.tensor.dtype)
+        new.tensor = dq
+        new.metadata["quantization_scale"] = scale.item()
+        return new
+
+
+# Alias matching the spec's example name; same fake-quant strategy, kept
+# distinct so a future PR can swap in a real fp8 kernel behind this class
+# without touching Int8QuantizationPass or any user code that references it.
+class FP8Pass(EAEPass):
+    name = "FP8Pass"
+
+    def __init__(self):
+        self._impl = Int8QuantizationPass()
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        return self._impl.process(adjoint, context)

eae_runtime/passes/regularization.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ..adjoint import AdjointState
+from .base import EAEPass
+
+
+class RegularizationPass(EAEPass):
+    """Adds a decay term proportional to the adjoint itself, i.e. a simple
+    stand-in for gradient-space regularization (e.g. an L2 penalty on the
+    *gradient*, distinct from weight decay which acts on parameters)."""
+
+    name = "RegularizationPass"
+
+    def __init__(self, strength: float = 0.0):
+        self.strength = strength
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        if self.strength == 0.0:
+            return adjoint.clone()
+        new = adjoint.clone()
+        new.tensor = new.tensor * (1.0 - self.strength)
+        return new
+
+
+class GaussianNoisePass(EAEPass):
+    """Injects zero-mean Gaussian noise into the adjoint, useful for studying
+    differential-privacy style gradient noising or robustness."""
+
+    name = "GaussianNoisePass"
+
+    def __init__(self, std: float = 0.0, generator: Optional[torch.Generator] = None):
+        self.std = std
+        self.generator = generator
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        if self.std == 0.0:
+            return adjoint.clone()
+        new = adjoint.clone()
+        noise = torch.randn(
+            new.tensor.shape, dtype=new.tensor.dtype, device=new.tensor.device, generator=self.generator
+        ) * self.std
+        new.tensor = new.tensor + noise
+        return new

eae_runtime/passes/synthetic_gradient.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from ..adjoint import AdjointState
+from .base import EAEPass
+
+
+class SyntheticGradientPass(EAEPass):
+    """Decoupled Neural Interfaces style synthetic gradient pass.
+
+    Wraps a small learned predictor `synthesizer(x) -> approx_grad` that can
+    be substituted for the true adjoint once it has been warmed up, letting
+    a researcher study update-unlocking / async credit assignment schemes.
+    During `warmup_steps`, the true adjoint passes through unchanged and is
+    used purely as a regression target to train the synthesizer; after
+    warmup, the pass optionally *replaces* the adjoint with the synthetic
+    estimate when `use_synthetic=True`.
+    """
+
+    name = "SyntheticGradientPass"
+
+    def __init__(
+        self,
+        feature_dim: int,
+        hidden_dim: int = 64,
+        lr: float = 1e-3,
+        warmup_steps: int = 0,
+        use_synthetic: bool = False,
+        device: Optional[torch.device] = None,
+    ):
+        self.synthesizer = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, feature_dim),
+        )
+        if device is not None:
+            self.synthesizer = self.synthesizer.to(device)
+        self.optimizer = torch.optim.Adam(self.synthesizer.parameters(), lr=lr)
+        self.warmup_steps = warmup_steps
+        self.use_synthetic = use_synthetic
+        self._step = 0
+        self.last_loss: Optional[float] = None
+
+    def process(self, adjoint: AdjointState, context=None) -> AdjointState:
+        new = adjoint.clone()
+        flat = new.tensor.reshape(new.tensor.shape[0], -1) if new.tensor.dim() > 1 else new.tensor.unsqueeze(0)
+
+        with torch.enable_grad():
+            pred = self.synthesizer(flat.detach().float())
+            target = flat.detach().float()
+            loss = torch.nn.functional.mse_loss(pred, target)
+            self.optimizer.zero_grad()
+            loss.backward()
+            self.optimizer.step()
+        self.last_loss = loss.item()
+        self._step += 1
+
+        if self.use_synthetic and self._step > self.warmup_steps:
+            with torch.no_grad():
+                synthetic = self.synthesizer(flat.float()).to(new.tensor.dtype)
+            new.tensor = synthetic.reshape(new.tensor.shape)
+            new.metadata["synthetic_gradient_loss"] = self.last_loss
+        return new

eae_runtime/pipeline.py

@@ -0,0 +1,58 @@
+"""
+AdjointPipeline: the central abstraction that turns EAE from a memory
+optimization into a research platform.
+
+Every reverse step exchanges an AdjointState through this programmable
+pipeline instead of the scheduler/reconstruction engine passing tensors
+directly between steps:
+
+    AdjointState -> Pass1 -> Pass2 -> ... -> PassN -> next local VJP
+
+A researcher working on gradient compression, synthetic gradients,
+error-feedback, distributed communication or adaptive optimization only
+ever implements a new Pass. They never need to understand or modify the
+reconstruction engine, scheduler, or memory manager.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import List, Optional
+
+from .adjoint import AdjointState
+from .events import EventBus, EventType
+from .passes.base import EAEPass
+
+
+class AdjointPipeline:
+    def __init__(self, passes: Optional[List[EAEPass]] = None, event_bus: Optional[EventBus] = None,
+                 profiler=None):
+        self.passes: List[EAEPass] = list(passes) if passes else []
+        self.event_bus = event_bus or EventBus()
+        self.profiler = profiler
+
+    def add_pass(self, p: EAEPass) -> None:
+        self.passes.append(p)
+
+    def run(self, adjoint: AdjointState, context: Optional[dict] = None) -> AdjointState:
+        context = context or {}
+        for p in self.passes:
+            start = time.perf_counter()
+            adjoint = p(adjoint, context)
+            elapsed = time.perf_counter() - start
+            if self.profiler is not None:
+                self.profiler.record("pass:" + p.name, elapsed)
+            self.event_bus.emit(EventType.PASS_APPLIED, pass_name=p.name, layer_id=adjoint.layer_id,
+                                 elapsed_seconds=elapsed)
+        return adjoint
+
+    def __call__(self, adjoint: AdjointState, context: Optional[dict] = None) -> AdjointState:
+        return self.run(adjoint, context)
+
+    def __len__(self) -> int:
+        return len(self.passes)
+
+
+# Backwards-compatible alias matching the spec's "Pass Manager" component
+# name; AdjointPipeline is the concrete implementation of that role.
+PassManager = AdjointPipeline

eae_runtime/profiler.py

@@ -0,0 +1,47 @@
+"""
+Built-in profiler: per-block reconstruction time, VJP time, pass time,
+memory, allocations, synchronization. Researchers need measurements.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from collections import defaultdict
+from typing import Dict, List
+
+
+class Profiler:
+    def __init__(self, enabled: bool = True):
+        self.enabled = enabled
+        self._records: Dict[str, List[float]] = defaultdict(list)
+
+    def record(self, name: str, elapsed_seconds: float) -> None:
+        if self.enabled:
+            self._records[name].append(elapsed_seconds)
+
+    @contextlib.contextmanager
+    def track(self, name: str):
+        if not self.enabled:
+            yield
+            return
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            self.record(name, time.perf_counter() - start)
+
+    def report(self) -> Dict[str, Dict[str, float]]:
+        out = {}
+        for name, values in self._records.items():
+            out[name] = {
+                "count": len(values),
+                "total_seconds": sum(values),
+                "mean_seconds": sum(values) / len(values) if values else 0.0,
+                "max_seconds": max(values) if values else 0.0,
+                "min_seconds": min(values) if values else 0.0,
+            }
+        return out
+
+    def reset(self) -> None:
+        self._records.clear()

eae_runtime/reconstruction.py

@@ -0,0 +1,89 @@
+"""
+Reconstruction Engine: given a block, a saved boundary activation, and an
+incoming AdjointState, it rebuilds the local computation graph, computes the
+VJP using PyTorch's *local* autograd (torch.autograd.grad), and returns a
+new AdjointState plus parameter gradients.
+
+No global graph ever exists. PyTorch owns local autograd / VJP / kernels;
+the runtime owns everything around it (see spec, Primary Design Principle 2).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from .adjoint import AdjointState
+from .events import EventBus, EventType
+
+
+class ReconstructionEngine:
+    def __init__(self, event_bus: Optional[EventBus] = None, compute_dtype: Optional[torch.dtype] = None):
+        self.event_bus = event_bus or EventBus()
+        self.compute_dtype = compute_dtype
+
+    def reconstruct(
+        self,
+        block: nn.Module,
+        input_activation: torch.Tensor,
+        adjoint: AdjointState,
+        block_name: str = "",
+    ) -> Tuple[AdjointState, Dict[nn.Parameter, torch.Tensor]]:
+        """Rebuild `block`'s local graph from `input_activation`, run the
+        forward again with autograd enabled, and back-propagate `adjoint`
+        through it using local VJP only.
+
+        Returns:
+            (new_adjoint_for_block_input, {param: grad_tensor})
+        """
+        compute_dtype = self.compute_dtype or input_activation.dtype
+
+        x = input_activation.detach().to(compute_dtype).clone().requires_grad_(True)
+
+        with torch.enable_grad():
+            out = block(x)
+
+        params = [p for p in block.parameters() if p.requires_grad]
+
+        grad_outputs = adjoint.tensor.to(device=out.device, dtype=out.dtype)
+        if grad_outputs.shape != out.shape:
+            raise RuntimeError(
+                f"Adjoint shape {tuple(grad_outputs.shape)} does not match "
+                f"block output shape {tuple(out.shape)} for block '{block_name}'"
+            )
+
+        inputs: List[torch.Tensor] = [x, *params]
+        grads = torch.autograd.grad(
+            outputs=out,
+            inputs=inputs,
+            grad_outputs=grad_outputs,
+            retain_graph=False,
+            allow_unused=True,
+        )
+        grad_x, param_grads_list = grads[0], grads[1:]
+
+        if grad_x is None:
+            grad_x = torch.zeros_like(x)
+
+        param_grads = {p: g for p, g in zip(params, param_grads_list) if g is not None}
+
+        new_adjoint = AdjointState(
+            tensor=grad_x.detach(),
+            layer_id=adjoint.layer_id - 1,
+            block=block_name,
+            metadata=dict(adjoint.metadata),
+        )
+        new_adjoint.record(f"reconstruct:{block_name}")
+
+        self.event_bus.emit(
+            EventType.BLOCK_RECONSTRUCTED,
+            block=block_name,
+            layer_id=adjoint.layer_id,
+            grad_norm=float(torch.linalg.vector_norm(grad_x.float()).item()),
+        )
+
+        # free the local graph explicitly
+        del out
+        return new_adjoint, param_grads

eae_runtime/runtime.py

@@ -0,0 +1,180 @@
+"""
+EAERuntime: the top-level object users instantiate. Ties together the
+Forward Executor, Boundary Store, Reverse Scheduler, Reconstruction Engine,
+Memory Manager, Adjoint Pipeline, Backend Manager, Event Bus and Profiler.
+
+    runtime = EAERuntime(model, optimizer, config)
+    runtime.train_step(batch)
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from .adjoint import AdjointState
+from .backend import BackendManager
+from .blocks import BlockDecomposer
+from .boundary_store import BoundaryStore
+from .config import RuntimeConfig
+from .events import EventBus, EventType
+from .forward_executor import ForwardExecutor
+from .logging_utils import attach_structured_logging
+from .memory import MemoryManager
+from .pipeline import AdjointPipeline
+from .profiler import Profiler
+from .reconstruction import ReconstructionEngine
+from .schedulers import ReverseContext, build_scheduler
+
+
+class EAERuntime:
+    def __init__(
+        self,
+        model: Union[nn.Sequential, List[nn.Module], nn.Module],
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        config: Optional[RuntimeConfig] = None,
+    ):
+        self.config = config or RuntimeConfig()
+        self.optimizer = optimizer
+
+        self.blocks: List[nn.Module] = BlockDecomposer.decompose(model)
+        self.block_names: List[str] = [
+            getattr(b, "eae_name", None) or f"{type(b).__name__}_{i}" for i, b in enumerate(self.blocks)
+        ]
+
+        self.event_bus = EventBus()
+        if self.config.log_level:
+            attach_structured_logging(self.event_bus, level=self.config.log_level)
+
+        self.profiler = Profiler(enabled=self.config.enable_profiler)
+
+        self.backend_manager = (
+            self.config.backend if isinstance(self.config.backend, BackendManager) else BackendManager(self.config.backend)
+        )
+
+        self.memory_manager = MemoryManager(policy=self.config.memory, event_bus=self.event_bus)
+
+        self.boundary_store = BoundaryStore(
+            offload=self.config.boundary_offload,
+            precision=self.config.boundary_precision,
+            pin_memory=self.config.pin_memory,
+        )
+
+        self.forward_executor = ForwardExecutor(event_bus=self.event_bus, compute_dtype=self.config.compute_dtype)
+        self.reconstruction_engine = ReconstructionEngine(event_bus=self.event_bus, compute_dtype=self.config.compute_dtype)
+
+        passes = list(self.config.passes)
+        self.pipeline = AdjointPipeline(passes=passes, event_bus=self.event_bus, profiler=self.profiler)
+
+        self.scheduler = build_scheduler(self.config.scheduler, num_microbatches=self.config.num_microbatches)
+
+        self.last_loss: Optional[float] = None
+        self.last_grad_stats: Dict[str, float] = {}
+
+    # ------------------------------------------------------------------ #
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run the detached forward pass, populating the boundary store, and
+        return the (detached) model output."""
+        self.boundary_store.clear()
+        with self.profiler.track("forward"):
+            out = self.forward_executor.run(self.blocks, x, self.boundary_store)
+        return out
+
+    def backward(self, output: torch.Tensor, loss_fn: Callable[[torch.Tensor], torch.Tensor]):
+        """Given the detached forward output, compute the loss (with a fresh
+        leaf tensor so we can extract dLoss/dOutput via ordinary autograd),
+        then run the reverse scheduler through the whole block stack.
+
+        Returns (loss_tensor, param_grads).
+        """
+        output_leaf = output.detach().clone().requires_grad_(True)
+        with torch.enable_grad():
+            loss = loss_fn(output_leaf)
+            loss.backward()
+
+        if output_leaf.grad is None:
+            raise RuntimeError(
+                "loss_fn produced no gradient w.r.t. the model output; make "
+                "sure it actually depends on its input tensor."
+            )
+
+        initial_adjoint = AdjointState(
+            tensor=output_leaf.grad.detach(),
+            layer_id=len(self.blocks),
+            block="output",
+        )
+        initial_adjoint = self.pipeline.run(initial_adjoint, context={"block_index": len(self.blocks), "block_name": "output"})
+
+        context = ReverseContext(
+            blocks=self.blocks,
+            block_names=self.block_names,
+            boundary_store=self.boundary_store,
+            reconstruction_engine=self.reconstruction_engine,
+            pipeline=self.pipeline,
+            memory_manager=self.memory_manager,
+            event_bus=self.event_bus,
+            profiler=self.profiler,
+            initial_adjoint=initial_adjoint,
+        )
+
+        with self.profiler.track("reverse_pass"):
+            param_grads = self.scheduler.run(context)
+
+        if self.config.grad_clip_norm is not None:
+            squared_norms = [(g.float() ** 2).sum() for g in param_grads.values()]
+            total_norm = torch.stack(squared_norms).sum().sqrt() if squared_norms else torch.tensor(0.0)
+            if total_norm > self.config.grad_clip_norm:
+                scale = self.config.grad_clip_norm / (total_norm + 1e-6)
+                param_grads = {p: g * scale for p, g in param_grads.items()}
+
+        return loss.detach(), param_grads
+
+    def apply_gradients(self, param_grads: Dict[nn.Parameter, torch.Tensor]) -> None:
+        for p, g in param_grads.items():
+            if p.grad is None:
+                p.grad = g.clone()
+            else:
+                p.grad = p.grad + g
+
+    def zero_grad(self) -> None:
+        for block in self.blocks:
+            for p in block.parameters():
+                p.grad = None
+
+    def train_step(self, x: torch.Tensor, loss_fn: Callable[[torch.Tensor], torch.Tensor]) -> float:
+        self.event_bus.emit(EventType.TRAIN_STEP_STARTED)
+        self.zero_grad()
+
+        out = self.forward(x)
+        loss, param_grads = self.backward(out, loss_fn)
+        self.apply_gradients(param_grads)
+
+        if self.optimizer is not None:
+            self.optimizer.step()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        self.last_loss = loss.item()
+        total_grad_norm = 0.0
+        if param_grads:
+            squared_norms = [(g.float() ** 2).sum() for g in param_grads.values()]
+            total_grad_norm = float(torch.stack(squared_norms).sum().sqrt().item())
+        self.last_grad_stats = {
+            "num_params_with_grad": len(param_grads),
+            "total_grad_norm": total_grad_norm,
+        }
+        self.event_bus.emit(EventType.TRAIN_STEP_FINISHED, loss=self.last_loss, **self.last_grad_stats)
+        return self.last_loss
+
+    # convenience for tests / advanced users needing raw grads without an
+    # optimizer step:
+    def compute_gradients(self, x: torch.Tensor, loss_fn: Callable[[torch.Tensor], torch.Tensor]):
+        self.zero_grad()
+        out = self.forward(x)
+        loss, param_grads = self.backward(out, loss_fn)
+        return loss, param_grads
+
+    def parameters(self):
+        for block in self.blocks:
+            yield from block.parameters()

eae_runtime/schedulers/__init__.py

@@ -0,0 +1,38 @@
+from .base import BaseScheduler, ReverseContext
+from .sequential import SequentialScheduler
+from .async_scheduler import AsyncScheduler
+from .pipeline_scheduler import PipelineScheduler
+from .distributed import DistributedScheduler
+
+_REGISTRY = {
+    "sequential": SequentialScheduler,
+    "async": AsyncScheduler,
+    "pipeline": PipelineScheduler,
+    "distributed": DistributedScheduler,
+}
+
+
+def build_scheduler(spec, **kwargs) -> BaseScheduler:
+    if isinstance(spec, BaseScheduler):
+        return spec
+    if isinstance(spec, str):
+        if spec not in _REGISTRY:
+            raise ValueError(f"Unknown scheduler '{spec}', choices: {list(_REGISTRY)}")
+        cls = _REGISTRY[spec]
+        if spec == "pipeline":
+            return cls(num_microbatches=kwargs.get("num_microbatches", 1))
+        return cls()
+    if isinstance(spec, type) and issubclass(spec, BaseScheduler):
+        return spec()
+    raise TypeError(f"Cannot build scheduler from {spec!r}")
+
+
+__all__ = [
+    "BaseScheduler",
+    "ReverseContext",
+    "SequentialScheduler",
+    "AsyncScheduler",
+    "PipelineScheduler",
+    "DistributedScheduler",
+    "build_scheduler",
+]

eae_runtime/schedulers/async_scheduler.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import concurrent.futures
+from typing import Dict
+
+import torch
+import torch.nn as nn
+
+from ..events import EventType
+from .base import BaseScheduler, ReverseContext
+
+
+class AsyncScheduler(BaseScheduler):
+    """Same numerical result as SequentialScheduler, but overlaps the (I/O
+    bound) boundary-activation fetch for block i-1 with the (compute bound)
+    reconstruction of block i using a background thread. This matters most
+    when the BoundaryStore is CPU-offloaded and the device is CUDA.
+
+    The reverse dependency chain (adjoint_{i-1} depends on adjoint_i) is
+    inherently sequential, so only the *fetch* is prefetched - the VJP
+    compute itself still runs in strict order, guaranteeing identical
+    results to SequentialScheduler.
+    """
+
+    name = "AsyncScheduler"
+
+    def __init__(self, max_workers: int = 2):
+        self.max_workers = max_workers
+
+    def run(self, context: ReverseContext) -> Dict[nn.Parameter, torch.Tensor]:
+        num_blocks = len(context.blocks)
+        order = list(reversed(range(num_blocks)))
+        adjoint = context.initial_adjoint
+        param_grads: Dict[nn.Parameter, torch.Tensor] = {}
+
+        if not order:
+            return param_grads
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            def fetch(i):
+                return context.boundary_store.get(i)
+
+            pending = {order[0]: executor.submit(fetch, order[0])}
+
+            for pos, i in enumerate(order):
+                input_activation = pending.pop(i).result()
+
+                if pos + 1 < len(order):
+                    nxt = order[pos + 1]
+                    pending[nxt] = executor.submit(fetch, nxt)
+
+                block = context.blocks[i]
+                name = context.block_names[i]
+                with context.profiler.track(f"reconstruct:{name}"):
+                    new_adjoint, updates = context.reconstruction_engine.reconstruct(
+                        block, input_activation, adjoint, block_name=name
+                    )
+                new_adjoint = context.pipeline.run(new_adjoint, context={"block_index": i, "block_name": name})
+                context.memory_manager.release(input_activation)
+                context.event_bus.emit(EventType.SCHEDULER_STEP, block_index=i, block_name=name, async_prefetch=True)
+
+                adjoint = new_adjoint
+                self._accumulate(param_grads, updates)
+
+        return param_grads