feat(ml): add backward_and_clip and apply_grads to Adam optimizer

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ml/src/lib.rs

@@ -49,6 +49,7 @@
 )]
 
 // Import common types properly - NO ALIASES THAT CONFLICT!
+use candle_core::backprop::GradStore;
 use candle_core::Tensor;
 use candle_core::Var;
 use candle_nn::Optimizer; // For Adam optimizer support
@@ -234,6 +235,64 @@ impl Adam {
         Ok(clipped_grad_norm)
     }
 
+    /// Perform backward pass with gradient clipping but WITHOUT an optimizer step.
+    ///
+    /// This is useful for gradient accumulation workflows where you want to
+    /// accumulate clipped gradients across multiple micro-batches before
+    /// applying a single optimizer step.
+    ///
+    /// # Arguments
+    ///
+    /// * `loss` - The loss tensor to compute gradients from
+    /// * `max_norm` - Maximum allowed gradient norm (gradients will be clipped to this value)
+    ///
+    /// # Returns
+    ///
+    /// Returns `Ok((grads, clipped_norm))` on success, where `grads` is the clipped
+    /// `GradStore` and `clipped_norm` is the gradient norm after clipping.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if:
+    /// - The backward pass fails to compute gradients
+    /// - Gradient clipping fails
+    pub fn backward_and_clip(
+        &self,
+        loss: &Tensor,
+        max_norm: f64,
+    ) -> Result<(GradStore, f64), MLError> {
+        // Compute gradients via backward pass
+        let mut grads = loss
+            .backward()
+            .map_err(|e| MLError::TrainingError(format!("Backward pass failed: {}", e)))?;
+
+        // Apply gradient clipping in-place
+        let (_actual_norm, clipped_norm) =
+            crate::gradient_utils::clip_grad_norm(&self.vars, &mut grads, max_norm)
+                .map_err(|e| MLError::TrainingError(format!("Gradient clipping failed: {}", e)))?;
+
+        Ok((grads, clipped_norm))
+    }
+
+    /// Apply pre-computed gradients to the optimizer.
+    ///
+    /// This is the companion to `backward_and_clip` for gradient accumulation
+    /// workflows. After accumulating and scaling gradients, call this method
+    /// to perform the optimizer step.
+    ///
+    /// # Arguments
+    ///
+    /// * `grads` - Pre-computed (and possibly accumulated/scaled) gradients
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if the optimizer step fails
+    pub fn apply_grads(&mut self, grads: &GradStore) -> Result<(), MLError> {
+        Optimizer::step(&mut self.optimizer, grads)
+            .map_err(|e| MLError::TrainingError(format!("Optimizer step failed: {}", e)))?;
+        Ok(())
+    }
+
     /// Compute the L2 norm of all gradients
     fn compute_gradient_norm(
         &self,