Make _sqrt_positive_part ONNX-exportable

Summary:
Replace boolean indexing and torch.is_grad_enabled() control flow in _sqrt_positive_part with a pure torch.where implementation. The old code used ret[positive_mask] = torch.sqrt(x[positive_mask]) which produces an incorrect ONNX Where/index_put node with mismatched broadcast shapes when the model is exported via torch.onnx.export.

The new implementation substitutes 1.0 for non-positive values before sqrt (avoiding infinite gradient at sqrt(0)) and masks the result back to 0, preserving the zero-subgradient-at-zero property.

Fixes https://github.com/facebookresearch/pytorch3d/issues/2020

Reviewed By: sgrigory

Differential Revision: D94365479

fbshipit-source-id: a1ebe8dc077573f83efc262520b6669159b83ef0

pytorch3d/csrc/utils/vec2.h

@@ -19,7 +19,7 @@ template <
         std::is_same<T, double>::value || std::is_same<T, float>::value>>
 struct vec2 {
   T x, y;
-  typedef T scalar_t;
+  using scalar_t = T;
   vec2(T x, T y) : x(x), y(y) {}
 };
 

pytorch3d/csrc/utils/vec3.h

@@ -18,7 +18,7 @@ template <
         std::is_same<T, double>::value || std::is_same<T, float>::value>>
 struct vec3 {
   T x, y, z;
-  typedef T scalar_t;
+  using scalar_t = T;
   vec3(T x, T y, T z) : x(x), y(y), z(z) {}
 };
 

pytorch3d/ops/cameras_alignment.py

@@ -223,8 +223,7 @@ def _align_camera_extrinsics(
         # of centered A and centered B
         Ac = A - Amu
         Bc = B - Bmu
-        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
-        align_t_s = (Ac * Bc).mean() / (Ac**2).mean().clamp(eps)
+        align_t_s = (Ac * Bc).mean() / torch.square(Ac).mean().clamp(eps)
     else:
         # set the scale to identity
         align_t_s = 1.0

pytorch3d/transforms/rotation_conversions.py

@@ -94,13 +94,9 @@ def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
     Returns torch.sqrt(torch.max(0, x))
     but with a zero subgradient where x is 0.
     """
-    ret = torch.zeros_like(x)
     positive_mask = x > 0
-    if torch.is_grad_enabled():
-        ret[positive_mask] = torch.sqrt(x[positive_mask])
-    else:
-        ret = torch.where(positive_mask, torch.sqrt(x), ret)
-    return ret
+    safe_x = torch.where(positive_mask, x, 1.0)
+    return torch.where(positive_mask, torch.sqrt(safe_x), 0.0)
 
 
 def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: