Use IMPLICIT for test_theseus_layer and fix related bugs

tests/test_theseus_layer.py

@@ -321,7 +321,16 @@ def cost_weight_fn():
             cost_weight_param_name: cost_weight_fn(),
         }
         pred_vars, info = layer_to_learn.forward(
-            input_values, optimizer_kwargs={**optimizer_kwargs, **{"verbose": verbose}}
+            input_values,
+            optimizer_kwargs={
+                **optimizer_kwargs,
+                **{
+                    "verbose": verbose,
+                    "backward_mode": "implicit"
+                    if learning_method == "direct"
+                    else "unroll",
+                },
+            },
         )
         assert not (
             (info.status == th.NonlinearOptimizerStatus.START)

theseus/optimizer/autograd/lu_cuda_sparse_autograd.py

@@ -51,6 +51,10 @@ def forward(  # type: ignore
             batch_size, A_row_ptr, A_col_ind, A_val_double, AtA_row_ptr, AtA_col_ind
         )
         if damping_alpha_beta is not None:
+            damping_alpha_beta = (
+                damping_alpha_beta[0].double(),
+                damping_alpha_beta[1].double(),
+            )
             AtA_args = sparse_structure.num_cols, AtA_row_ptr, AtA_col_ind, AtA
             apply_damping(batch_size, *AtA_args, *damping_alpha_beta)
         solver_context.factor(AtA)

theseus/optimizer/nonlinear/nonlinear_optimizer.py

@@ -378,6 +378,7 @@ def _optimize_loop(
                 info.last_err,
                 converged_indices,
                 force_update,
+                truncated_grad_loop=truncated_grad_loop,
                 **kwargs,
             )  # err is shape (batch_size,)
             if all_rejected:
@@ -566,12 +567,19 @@ def _step(
         previous_err: torch.Tensor,
         converged_indices: torch.Tensor,
         force_update: bool,
+        truncated_grad_loop: bool,
         **kwargs,
     ) -> Tuple[torch.Tensor, bool]:
         tensor_dict, err = self._compute_retracted_tensors_and_error(
             delta, converged_indices, force_update
         )
-        reject_indices = self._complete_step(delta, err, previous_err, **kwargs)
+        if truncated_grad_loop:
+            # For "implicit" or "truncated", the grad-attached steps are just GN steps
+            # So, we need to avoid calling `_complete_step`, as it's likely to reject
+            # the step computed
+            reject_indices = None
+        else:
+            reject_indices = self._complete_step(delta, err, previous_err, **kwargs)
 
         if reject_indices is not None and reject_indices.all():
             return previous_err, True