Preserve FP32 coordinates under BF16 in tomography and for INRs

src/quantem/core/ml/models/kplanes.py

@@ -309,6 +309,13 @@ def resolution(self, resolution: Sequence[int]):
 # ---------------------------------------------------------------------------
 
 
+def _rotate_tilted_points_fp32(
+    pts: torch.Tensor, rotation_matrices: torch.Tensor
+) -> torch.Tensor:
+    with torch.autocast(device_type=pts.device.type, enabled=False):
+        return torch.einsum("tij,bj->tbi", rotation_matrices.float(), pts.float())
+
+
 def interpolate_ms_features_tilted(
     pts: torch.Tensor,  # (B, 3)
     ms_grids: nn.ParameterList,  # each grid: (3*T, C, H, W)
@@ -322,7 +329,7 @@ def interpolate_ms_features_tilted(
     B = pts.shape[0]
 
     # (T, B, 3)  — rotate all points by all rotations at once
-    rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
+    rotated = _rotate_tilted_points_fp32(pts, rotation_matrices)
 
     # Build (T, 3, B, 2) coords for planes XY, ZX, YZ in one shot.
     # index_select is faster and cleaner than advanced indexing with python lists.
@@ -610,7 +617,7 @@ def interpolate_ms_features_cp_tilted(
     B = pts.shape[0]
 
     # Rotate all points by all rotations: (T, B, 3)
-    rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
+    rotated = _rotate_tilted_points_fp32(pts, rotation_matrices)
 
     per_scale_features = []
     for line_coef in ms_grids:

src/quantem/core/ml/models/so3params.py

@@ -174,11 +174,13 @@ def rotmat_to_r9(R: torch.Tensor) -> torch.Tensor:
     @staticmethod
     def r9_to_rotmat(M: torch.Tensor) -> torch.Tensor:
         """R9 (..., 3, 3) -> nearest SO(3) matrix via SVD+."""
-        U, _, Vh = torch.linalg.svd(M)
-        d = torch.det(U @ Vh)
-        diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M.dtype)
-        diag[..., 2] = d
-        return U @ (diag.unsqueeze(-1) * Vh)
+        with torch.autocast(device_type=M.device.type, enabled=False):
+            M_work = M if M.dtype in (torch.float32, torch.float64) else M.float()
+            U, _, Vh = torch.linalg.svd(M_work)
+            d = torch.det(U @ Vh)
+            diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M_work.dtype)
+            diag[..., 2] = d
+            return U @ (diag.unsqueeze(-1) * Vh)
 
     def as_matrix(self) -> torch.Tensor:
         return self.r9_to_rotmat(self.M)
@@ -189,4 +191,4 @@ def from_matrix(cls, R: torch.Tensor) -> "SO3ParamR9SVD":
         obj = cls(R.shape[0], init="identity")
         with torch.no_grad():
             obj.M.copy_(cls.rotmat_to_r9(R))
-        return obj
+        return obj

src/quantem/tomography/dataset_models.py

@@ -508,32 +508,35 @@ def forward(self, dummy_input: Any = None):
     def get_coords(
         self, batch: dict[str, torch.Tensor], N: int, num_samples_per_ray: int
     ) -> torch.Tensor:
-        pixel_i = batch["pixel_i"].float().to(self.device, non_blocking=True)
-        pixel_j = batch["pixel_j"].float().to(self.device, non_blocking=True)
-        # target_values = batch["target_value"].to(self.device, non_blocking=True)
-        phis = batch["phi"].to(self.device, non_blocking=True)
-        projection_indices = batch["projection_idx"].to(self.device, non_blocking=True)
-        with torch.no_grad():
-            batch_ray_coords = self.create_batch_rays(pixel_i, pixel_j, N, num_samples_per_ray)
-
-        shifts, z1_params, z3_params = self.forward(None)
-        batch_shifts = torch.index_select(shifts, 0, projection_indices)
-        batch_z1 = torch.index_select(z1_params, 0, projection_indices)
-        batch_z3 = torch.index_select(z3_params, 0, projection_indices)
-
-        transformed_rays = self.transform_batch_rays(
-            batch_ray_coords,
-            z1=batch_z1,
-            x=phis,
-            z3=batch_z3,
-            shifts=batch_shifts,
-            N=N,
-            sampling_rate=1.0,
-        )
-        all_coords = transformed_rays.view(-1, 3)
-
-        all_coords = all_coords.to(self.device, dtype=torch.float32, non_blocking=True)
-        return all_coords
+        with torch.autocast(device_type=self.device.type, enabled=False):
+            pixel_i = batch["pixel_i"].to(
+                self.device, dtype=torch.float32, non_blocking=True
+            )
+            pixel_j = batch["pixel_j"].to(
+                self.device, dtype=torch.float32, non_blocking=True
+            )
+            phis = batch["phi"].to(self.device, dtype=torch.float32, non_blocking=True)
+            projection_indices = batch["projection_idx"].to(self.device, non_blocking=True)
+            with torch.no_grad():
+                batch_ray_coords = self.create_batch_rays(
+                    pixel_i, pixel_j, N, num_samples_per_ray
+                )
+
+            shifts, z1_params, z3_params = self.forward(None)
+            batch_shifts = torch.index_select(shifts.float(), 0, projection_indices)
+            batch_z1 = torch.index_select(z1_params.float(), 0, projection_indices)
+            batch_z3 = torch.index_select(z3_params.float(), 0, projection_indices)
+
+            transformed_rays = self.transform_batch_rays(
+                batch_ray_coords,
+                z1=batch_z1,
+                x=phis,
+                z3=batch_z3,
+                shifts=batch_shifts,
+                N=N,
+                sampling_rate=1.0,
+            )
+            return transformed_rays.reshape(-1, 3)
 
     @staticmethod
     def create_batch_rays(
@@ -542,9 +545,17 @@ def create_batch_rays(
         batch_size = len(pixel_i)
         x_coords = (pixel_j / (N - 1)) * 2 - 1
         y_coords = (pixel_i / (N - 1)) * 2 - 1
-        z_coords = torch.linspace(-1, 1, num_samples_per_ray, device=pixel_i.device)
+        z_coords = torch.linspace(
+            -1, 1, num_samples_per_ray, device=pixel_i.device, dtype=torch.float32
+        )
 
-        rays = torch.zeros(batch_size, num_samples_per_ray, 3, device=pixel_i.device)
+        rays = torch.zeros(
+            batch_size,
+            num_samples_per_ray,
+            3,
+            device=pixel_i.device,
+            dtype=torch.float32,
+        )
 
         rays[:, :, 0] = x_coords.unsqueeze(1)
         rays[:, :, 1] = y_coords.unsqueeze(1)

src/quantem/tomography/object_models.py

@@ -20,6 +20,14 @@
 from quantem.tomography.tomography_context import ReconstructionContext
 
 
+def _tomography_autocast(device: torch.device | str, use_bfloat16: bool):
+    device = torch.device(device)
+    enabled = bool(
+        use_bfloat16 and device.type == "cuda" and torch.cuda.is_bf16_supported()
+    )
+    return torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=enabled)
+
+
 class ObjConstraintParams:
     """
     Namespace class for object reconstruction constraint dataclasses and parsing utilities.
@@ -687,6 +695,7 @@ def pretrain(
         scheduler_params: dict | None = None,
         loss_fn: Callable | str = "l1",
         verbose: bool = True,
+        use_bfloat16: bool = True,
     ):
         """
         Pretrain the INR model to fit target volume.
@@ -717,13 +726,15 @@ def pretrain(
             num_iters=num_iters,
             loss_fn=loss_fn,
             verbose=verbose,
+            use_bfloat16=use_bfloat16,
         )
 
     def _pretrain(
         self,
         num_iters: int,
         loss_fn: Callable,
         verbose: bool,
+        use_bfloat16: bool,
     ):
         if self.optimizer is None:
             raise RuntimeError("Optimizer not set. Call set_optimizer() first.")
@@ -741,9 +752,7 @@ def _pretrain(
                 coords = batch["coords"].to(self.device, non_blocking=True)
                 target = batch["target"].to(self.device, non_blocking=True)
 
-                with torch.autocast(
-                    device_type=self.device.type, dtype=torch.bfloat16, enabled=True
-                ):
+                with _tomography_autocast(self.device, use_bfloat16):
                     outputs = self.forward(coords)
                     loss = loss_fn(outputs, target)
 

src/quantem/tomography/tomography.py

@@ -27,6 +27,7 @@
     ObjectINR,
     ObjectPixelated,
     ObjectTensorDecomp,
+    _tomography_autocast,
 )
 from quantem.tomography.radon.radon import iradon_torch, radon_torch
 from quantem.tomography.tomography_base import TomographyBase
@@ -84,6 +85,7 @@ def reconstruct(
         loss_func_kwargs: dict = {},
         reset_dset: DatasetModelType | None = None,
         show_metrics: bool = False,
+        use_bfloat16: bool = True,
     ):
         """
         This function should be able to handle both AD and INR-based tomography reconstruction methods.
@@ -203,11 +205,7 @@ def reconstruct(
 
             for batch_idx, batch in enumerate(self.dataloader):
                 self.zero_grad_all()
-                with torch.autocast(
-                    device_type=self.device.type,
-                    dtype=torch.bfloat16,
-                    enabled=False,
-                ):
+                with _tomography_autocast(self.device, use_bfloat16):
                     all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
                     all_densities = self.obj_model.forward(all_coords)
@@ -285,11 +283,7 @@ def reconstruct(
                     val_loss = torch.tensor(0.0, device=self.device)
 
                     for batch in self.val_dataloader:
-                        with torch.autocast(
-                            device_type=self.device.type,
-                            dtype=torch.bfloat16,
-                            enabled=True,
-                        ):
+                        with _tomography_autocast(self.device, use_bfloat16):
                             all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
 
                             all_densities = self.obj_model.forward(all_coords)