Fix an issue where if radii is not provided (it is optional), a 'tensor does not have a device' error is thrown when 'save_for_backward' is called

Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>

Author: swahtz

src/fvdb/detail/autograd/EvaluateSphericalHarmonics.cpp

@@ -20,19 +20,27 @@ EvaluateSphericalHarmonics::forward(
         viewDirections,                                                   // [C, N, 3] (optional)
     const EvaluateSphericalHarmonics::Variable &sh0Coeffs,                // [N, 1, D]
     const std::optional<EvaluateSphericalHarmonics::Variable> &shNCoeffs, // [N, K-1, D]
-    const EvaluateSphericalHarmonics::Variable &radii                     // [C, N]
+    const std::optional<EvaluateSphericalHarmonics::Variable> &radii      // [C, N] (optional)
 ) {
     FVDB_FUNC_RANGE_WITH_NAME("EvaluateSphericalHarmonics::forward");
     const Variable viewDirectionsValue = viewDirections.value_or(torch::Tensor());
     const Variable shNCoeffsValue      = shNCoeffs.value_or(torch::Tensor());
+    const Variable radiiValue          = radii.value_or(torch::Tensor());
     const Variable renderQuantities    = FVDB_DISPATCH_KERNEL(sh0Coeffs.device(), [&]() {
         return ops::dispatchSphericalHarmonicsForward<DeviceTag>(
-            shDegreeToUse, numCameras, viewDirectionsValue, sh0Coeffs, shNCoeffsValue, radii);
+            shDegreeToUse, numCameras, viewDirectionsValue, sh0Coeffs, shNCoeffsValue, radiiValue);
     });
-    ctx->save_for_backward({viewDirectionsValue, shNCoeffsValue, radii});
+    // only save radii if defined to avoid device access issues
+    const bool hasRadii = radii.has_value() && radii.value().defined();
+    if (hasRadii) {
+        ctx->save_for_backward({viewDirectionsValue, shNCoeffsValue, radiiValue});
+    } else {
+        ctx->save_for_backward({viewDirectionsValue, shNCoeffsValue});
+    }
     ctx->saved_data["shDegreeToUse"] = static_cast<int64_t>(shDegreeToUse);
     ctx->saved_data["numCameras"]    = static_cast<int64_t>(numCameras);
     ctx->saved_data["numGaussians"]  = static_cast<int64_t>(sh0Coeffs.size(0));
+    ctx->saved_data["hasRadii"]      = hasRadii;
     return {renderQuantities};
 }
 
@@ -49,13 +57,16 @@ EvaluateSphericalHarmonics::backward(EvaluateSphericalHarmonics::AutogradContext
     VariableList saved = ctx->get_saved_variables();
     Variable viewDirs  = saved.at(0);
     Variable shNCoeffs = saved.at(1);
-    Variable radii     = saved.at(2);
 
     const int shDegreeToUse          = static_cast<int>(ctx->saved_data["shDegreeToUse"].toInt());
     const int numCameras             = static_cast<int>(ctx->saved_data["numCameras"].toInt());
     const int numGaussians           = static_cast<int>(ctx->saved_data["numGaussians"].toInt());
+    const bool hasRadii              = ctx->saved_data["hasRadii"].toBool();
     const bool computeDLossDViewDirs = ctx->needs_input_grad(1);
 
+    // Radii is only saved if it was defined in forward
+    Variable radii = hasRadii ? saved.at(2) : torch::Tensor();
+
     auto variables           = FVDB_DISPATCH_KERNEL(dLossdColors.device(), [&]() {
         return ops::dispatchSphericalHarmonicsBackward<DeviceTag>(shDegreeToUse,
                                                                   numCameras,

src/fvdb/detail/autograd/EvaluateSphericalHarmonics.h

@@ -22,7 +22,7 @@ struct EvaluateSphericalHarmonics : public torch::autograd::Function<EvaluateSph
             const std::optional<Variable> viewDirections, // [N, 3] or empty for deg 0
             const Variable &sh0Coeffs,                    // [N, 1, D]
             const std::optional<Variable> &shNCoeffs,     // [N, K-1, D]
-            const Variable &radii                         // [N,]
+            const std::optional<Variable> &radii          // [C, N] (optional)
     );
 
     static VariableList backward(AutogradContext *ctx, VariableList gradOutput);

src/python/GaussianSplatBinding.cpp

@@ -402,9 +402,8 @@ bind_gaussian_splat3d(py::module &m) {
            const std::optional<torch::Tensor> &shN,
            const std::optional<torch::Tensor> &viewDirections,
            const std::optional<torch::Tensor> &radii) {
-            torch::Tensor radiiValue = radii.value_or(torch::Tensor());
             return fvdb::detail::autograd::EvaluateSphericalHarmonics::apply(
-                shDegree, numCameras, viewDirections, sh0, shN, radiiValue)[0];
+                shDegree, numCameras, viewDirections, sh0, shN, radii)[0];
         },
         R"doc(
 Evaluate spherical harmonics to compute view-dependent features/colors.

tests/unit/test_gaussian_splat_3d.py

@@ -3138,6 +3138,103 @@ def test_degree_0_basic(self):
         # since there's no view dependence
         self.assertTrue(torch.allclose(result[0], result[1], atol=1e-6))
 
+    def test_optional_parameters_omitted(self):
+        """Test that optional parameters (shN, view_directions, radii) can be omitted.
+
+        This is a regression test for an issue where passing undefined tensors
+        to save_for_backward caused 'tensor does not have a device' errors.
+        """
+        N = 50
+        D = 3
+        C = 2
+
+        sh0 = torch.randn(N, 1, D, device=self.device)
+
+        # Test degree 0 with no optional parameters
+        result = evaluate_spherical_harmonics(
+            sh_degree=0,
+            num_cameras=C,
+            sh0=sh0,
+        )
+        self.assertEqual(result.shape, (C, N, D))
+        self.assertFalse(torch.isnan(result).any())
+
+        # Test degree 0 with explicit None for optional parameters
+        result = evaluate_spherical_harmonics(
+            sh_degree=0,
+            num_cameras=C,
+            sh0=sh0,
+            shN=None,
+            view_directions=None,
+            radii=None,
+        )
+        self.assertEqual(result.shape, (C, N, D))
+        self.assertFalse(torch.isnan(result).any())
+
+        # Test higher degree without radii
+        shN = torch.randn(N, 15, D, device=self.device)
+        view_dirs = torch.randn(C, N, 3, device=self.device)
+
+        result = evaluate_spherical_harmonics(
+            sh_degree=3,
+            num_cameras=C,
+            sh0=sh0,
+            shN=shN,
+            view_directions=view_dirs,
+            # radii intentionally omitted
+        )
+        self.assertEqual(result.shape, (C, N, D))
+        self.assertFalse(torch.isnan(result).any())
+
+        # Test higher degree with explicit radii=None
+        result = evaluate_spherical_harmonics(
+            sh_degree=3,
+            num_cameras=C,
+            sh0=sh0,
+            shN=shN,
+            view_directions=view_dirs,
+            radii=None,
+        )
+        self.assertEqual(result.shape, (C, N, D))
+        self.assertFalse(torch.isnan(result).any())
+
+    def test_optional_parameters_gradient_flow(self):
+        """Test gradient flow works when optional parameters are omitted.
+
+        Regression test to ensure backward pass works without radii.
+        """
+        N = 20
+        D = 3
+        C = 2
+
+        sh0 = torch.randn(N, 1, D, device=self.device, requires_grad=True)
+        shN = torch.randn(N, 15, D, device=self.device, requires_grad=True)
+        view_dirs = torch.randn(C, N, 3, device=self.device)
+
+        # Forward without radii
+        result = evaluate_spherical_harmonics(
+            sh_degree=3,
+            num_cameras=C,
+            sh0=sh0,
+            shN=shN,
+            view_directions=view_dirs,
+            # radii intentionally omitted
+        )
+
+        # Backward
+        loss = result.sum()
+        loss.backward()
+
+        # Check gradients exist and are valid
+        self.assertIsNotNone(sh0.grad)
+        self.assertIsNotNone(shN.grad)
+        assert sh0.grad is not None  # for type narrowing
+        assert shN.grad is not None  # for type narrowing
+        self.assertTrue(torch.any(sh0.grad != 0))
+        self.assertTrue(torch.any(shN.grad != 0))
+        self.assertFalse(torch.isnan(sh0.grad).any())
+        self.assertFalse(torch.isnan(shN.grad).any())
+
     def test_degree_0_matches_expected(self):
         """Test that degree 0 SH evaluation produces expected output."""
         N = 10