Read implicitron's own annotations via inspect.get_annotations

Summary:
`expand_args_fields` gated its whole member-processing pass on
`"__annotations__" in some_class.__dict__`. Under PEP 649 that entry is not in the
class dict until something materializes it, so on 3.14 the check was simply False,
no member was processed, and the pass that moves defaultable members to the end of
`__annotations__` never ran. The class then reached
`dataclasses.dataclass(eq=False)` with a defaulted `<name>_class_type` field ahead
of the non-defaulted `<name>` it replaces, which is a hard `TypeError` at class
construction.

`inspect.get_annotations` returns a class's own annotations — not a base's — which
is exactly what the `__dict__` lookup was expressing, and it behaves identically on
3.12. The in-place `del` / re-add on `some_class.__annotations__` further down is
unaffected: attribute access materializes and caches the dict on 3.14, so the
mutations stick and `dataclasses` sees them (verified on both interpreters).

325 canary failures in V58.

Reviewed By: bottler, ambv

Differential Revision: D115343017

fbshipit-source-id: 33f38dfac9f10c46b1f6b0e57f7005c9d7750dfa

.github/workflows/build.yml

@@ -21,3 +21,45 @@ jobs:
       run: |-
         conda create --name env --yes --quiet conda-build
         conda run --no-capture-output --name env python3 ./packaging/build_conda.py --use-conda-cuda
+
+  # Build-only verification for the ROCm/HIP code paths. Runs in an AMD ROCm
+  # dev container on a CPU-only GitHub runner; we don't need an AMD GPU just
+  # to compile, and not running tests keeps the CI cost low. Catches build
+  # regressions in the ROCm code paths (USE_ROCM guards, hipify-touched sources,
+  # the pulsar HIP intrinsic replacements, etc.).
+  linux_rocm_build:
+    runs-on: ubuntu-latest
+    container:
+      # `-complete` tag bundles the full ROCm math stack (rocThrust, hipCUB,
+      # rocPRIM, ...). The plain `7.2.3` tag is HIP-runtime-only and fails to
+      # find <thrust/complex.h> when including PyTorch headers.
+      image: rocm/dev-ubuntu-22.04:7.2.3-complete
+    env:
+      PYTORCH_VERSION: "2.11.0"
+      ROCM_INDEX: "rocm7.2"
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install Python and torch+rocm
+      run: |-
+        apt-get update
+        apt-get install -y --no-install-recommends python3 python3-dev python3-pip git
+        python3 -m pip install --upgrade pip
+        python3 -m pip install --index-url https://download.pytorch.org/whl/${ROCM_INDEX} torch==${PYTORCH_VERSION}
+    - name: Verify torch is ROCm-built
+      run: |-
+        python3 -c "import torch; assert torch.version.hip is not None, 'torch is not HIP-built'; print('torch.version.hip:', torch.version.hip)"
+    - name: Build pytorch3d _C extension (build only, no tests)
+      env:
+        # CPU-only runner: torch.cuda.is_available() is False, so force the
+        # CUDAExtension path. ROCM_HOME is auto-detected from /opt/rocm in
+        # the rocm/dev-ubuntu container.
+        FORCE_CUDA: "1"
+      run: |-
+        python3 -m pip install --no-build-isolation -v .
+    - name: Smoke import
+      # cd out of the checkout root so the source-tree pytorch3d/ directory
+      # (which has no _C.so since the build doesn't install in-place) doesn't
+      # shadow the site-packages install via sys.path[0] for `python -c`.
+      run: |-
+        cd /tmp
+        python3 -c "import torch; from pytorch3d import _C; print('PulsarRenderer:', hasattr(_C, 'PulsarRenderer')); print('n_symbols:', len([s for s in dir(_C) if not s.startswith('_')]))"

.gitignore

@@ -6,6 +6,17 @@ dist/
 **/.ipynb_checkpoints
 **/.ipynb_checkpoints/**
 
+# Build artifacts produced in-place by torch.utils.cpp_extension auto-hipify
+# when pytorch3d is built against a ROCm PyTorch.
+*.so
+pytorch3d/csrc/**/*.hip
+pytorch3d/csrc/**/*_hip.cpp
+pytorch3d/csrc/**/*_hip.h
+pytorch3d/csrc/**/*_hip.cuh
+
+# Debug PNG dumps written by pulsar tests when FB_TEST is unset.
+tests/pulsar/test_out/
+
 
 # Docusaurus site
 website/yarn.lock

dev/linter.sh

@@ -10,7 +10,7 @@
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 DIR=$(dirname "${DIR}")
 
-if [[ -f "${DIR}/TARGETS" ]]
+if [[ -f "${DIR}/BUCK" ]]
 then
   pyfmt "${DIR}"
 else

docs/conf.py

@@ -19,7 +19,6 @@
 #
 import os
 import sys
-
 import unittest.mock as mock
 
 from recommonmark.parser import CommonMarkParser

docs/modules/implicitron/datasets.rst

@@ -3,11 +3,6 @@ pytorch3d.implicitron.dataset specific datasets
 
 specific datasets
 
-.. automodule:: pytorch3d.implicitron.dataset.blender_dataset_map_provider
-    :members:
-    :undoc-members:
-    :show-inheritance:
-
 .. automodule:: pytorch3d.implicitron.dataset.json_index_dataset_map_provider
     :members:
     :undoc-members:
@@ -18,11 +13,6 @@ specific datasets
     :undoc-members:
     :show-inheritance:
 
-.. automodule:: pytorch3d.implicitron.dataset.llff_dataset_map_provider
-    :members:
-    :undoc-members:
-    :show-inheritance:
-
 .. automodule:: pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider
     :members:
     :undoc-members:

projects/implicitron_trainer/configs/overfit_singleseq_nerf_blender.yaml

@@ -1,56 +0,0 @@
-defaults:
-- overfit_singleseq_base
-- _self_
-exp_dir: "./data/overfit_nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
-data_source_ImplicitronDataSource_args:
-  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
-    dataset_length_train: 100
-  dataset_map_provider_class_type: BlenderDatasetMapProvider
-  dataset_map_provider_BlenderDatasetMapProvider_args:
-    base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
-    n_known_frames_for_test: null
-    object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
-    path_manager_factory_class_type: PathManagerFactory
-    path_manager_factory_PathManagerFactory_args:
-      silence_logs: true
-
-model_factory_ImplicitronModelFactory_args:
-  model_class_type: "OverfitModel"
-  model_OverfitModel_args:
-    mask_images: false
-    raysampler_class_type: AdaptiveRaySampler
-    raysampler_AdaptiveRaySampler_args:
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      n_rays_per_image_sampled_from_mask: 4096
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-      scene_extent: 2.0
-      scene_center:
-      - 0.0
-      - 0.0
-      - 0.0
-    renderer_MultiPassEmissionAbsorptionRenderer_args:
-      density_noise_std_train: 0.0
-      n_pts_per_ray_fine_training: 128
-      n_pts_per_ray_fine_evaluation: 128
-      raymarcher_EmissionAbsorptionRaymarcher_args:
-        blend_output: false
-    loss_weights:
-      loss_rgb_mse: 1.0
-      loss_prev_stage_rgb_mse: 1.0
-      loss_mask_bce: 0.0
-      loss_prev_stage_mask_bce: 0.0
-      loss_autodecoder_norm: 0.00
-
-optimizer_factory_ImplicitronOptimizerFactory_args:
-  exponential_lr_step_size: 3001
-  lr_policy: LinearExponential
-  linear_exponential_lr_milestone: 200
-
-training_loop_ImplicitronTrainingLoop_args:
-  max_epochs: 6000
-  metric_print_interval: 10
-  store_checkpoints_purge: 3
-  test_when_finished: true
-  validation_interval: 100

projects/implicitron_trainer/configs/repro_singleseq_nerf_blender.yaml

@@ -1,55 +0,0 @@
-defaults:
-- repro_singleseq_base
-- _self_
-exp_dir: "./data/nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
-data_source_ImplicitronDataSource_args:
-  data_loader_map_provider_SequenceDataLoaderMapProvider_args:
-    dataset_length_train: 100
-  dataset_map_provider_class_type: BlenderDatasetMapProvider
-  dataset_map_provider_BlenderDatasetMapProvider_args:
-    base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
-    n_known_frames_for_test: null
-    object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
-    path_manager_factory_class_type: PathManagerFactory
-    path_manager_factory_PathManagerFactory_args:
-      silence_logs: true
-
-model_factory_ImplicitronModelFactory_args:
-  model_GenericModel_args:
-    mask_images: false
-    raysampler_class_type: AdaptiveRaySampler
-    raysampler_AdaptiveRaySampler_args:
-      n_pts_per_ray_training: 64
-      n_pts_per_ray_evaluation: 64
-      n_rays_per_image_sampled_from_mask: 4096
-      stratified_point_sampling_training: true
-      stratified_point_sampling_evaluation: false
-      scene_extent: 2.0
-      scene_center:
-      - 0.0
-      - 0.0
-      - 0.0
-    renderer_MultiPassEmissionAbsorptionRenderer_args:
-      density_noise_std_train: 0.0
-      n_pts_per_ray_fine_training: 128
-      n_pts_per_ray_fine_evaluation: 128
-      raymarcher_EmissionAbsorptionRaymarcher_args:
-        blend_output: false
-    loss_weights:
-      loss_rgb_mse: 1.0
-      loss_prev_stage_rgb_mse: 1.0
-      loss_mask_bce: 0.0
-      loss_prev_stage_mask_bce: 0.0
-      loss_autodecoder_norm: 0.00
-
-optimizer_factory_ImplicitronOptimizerFactory_args:
-  exponential_lr_step_size: 3001
-  lr_policy: LinearExponential
-  linear_exponential_lr_milestone: 200
-
-training_loop_ImplicitronTrainingLoop_args:
-  max_epochs: 6000
-  metric_print_interval: 10
-  store_checkpoints_purge: 3
-  test_when_finished: true
-  validation_interval: 100

projects/implicitron_trainer/experiment.py

@@ -48,22 +48,18 @@ The outputs of the experiment are saved and logged in multiple ways:
 import logging
 import os
 import warnings
-
 from dataclasses import field
 
 import hydra
-
 import torch
 from accelerate import Accelerator
 from omegaconf import DictConfig, OmegaConf
 from packaging import version
-
 from pytorch3d.implicitron.dataset.data_source import (
     DataSourceBase,
     ImplicitronDataSource,
 )
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
-
 from pytorch3d.implicitron.models.renderer.multipass_ea import (
     MultiPassEmissionAbsorptionRenderer,
 )
@@ -197,6 +193,7 @@ class Experiment(Configurable):
             last_epoch=start_epoch,
             model=model,
             resume=self.model_factory.resume,
+            # pyrefly: ignore [missing-attribute]
             resume_epoch=self.model_factory.resume_epoch,
         )
 
@@ -216,6 +213,7 @@ class Experiment(Configurable):
 
         # Enter the main training loop.
         self.training_loop.run(
+            # pyrefly: ignore [bad-argument-type]
             train_loader=train_loader,
             val_loader=val_loader,
             test_loader=test_loader,

projects/implicitron_trainer/impl/model_factory.py

@@ -11,7 +11,6 @@ import os
 from typing import Optional
 
 import torch.optim
-
 from accelerate import Accelerator
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
 from pytorch3d.implicitron.tools import model_io

projects/implicitron_trainer/impl/optimizer_factory.py

@@ -14,9 +14,7 @@ from dataclasses import field
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch.optim
-
 from accelerate import Accelerator
-
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
 from pytorch3d.implicitron.tools import model_io
 from pytorch3d.implicitron.tools.config import (
@@ -171,9 +169,11 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
                 gamma=self.gamma,
             )
         elif self.lr_policy.casefold() == "Exponential".casefold():
+            # pyre-fixme[28]: Unexpected keyword argument `verbose`.
             scheduler = torch.optim.lr_scheduler.LambdaLR(
                 optimizer,
                 lambda epoch: self.gamma ** (epoch / self.exponential_lr_step_size),
+                # pyrefly: ignore [unexpected-keyword]
                 verbose=False,
             )
         elif self.lr_policy.casefold() == "LinearExponential".casefold():
@@ -190,8 +190,13 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
                     gamma = self.gamma ** (epoch_rest / self.exponential_lr_step_size)
                 return gamma
 
+            # pyre-fixme[28]: Unexpected keyword argument `verbose`.
             scheduler = torch.optim.lr_scheduler.LambdaLR(
-                optimizer, _get_lr, verbose=False
+                # pyrefly: ignore [unexpected-keyword]
+                optimizer,
+                _get_lr,
+                # pyrefly: ignore [unexpected-keyword]
+                verbose=False,
             )
         else:
             raise ValueError("no such lr policy %s" % self.lr_policy)

projects/implicitron_trainer/impl/training_loop.py

@@ -199,6 +199,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
                     and self.test_interval > 0
                     and epoch % self.test_interval == 0
                 ):
+                    # pyrefly: ignore [missing-attribute]
                     self.evaluator.run(
                         device=device,
                         dataloader=test_loader,
@@ -215,6 +216,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
 
         if self.test_when_finished:
             if test_loader is not None:
+                # pyrefly: ignore [missing-attribute]
                 self.evaluator.run(
                     device=device,
                     dump_to_json=True,

projects/implicitron_trainer/tests/experiment.yaml

@@ -13,13 +13,6 @@ hydra:
 data_source_ImplicitronDataSource_args:
   dataset_map_provider_class_type: ???
   data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
-  dataset_map_provider_BlenderDatasetMapProvider_args:
-    base_dir: ???
-    object_name: ???
-    path_manager_factory_class_type: PathManagerFactory
-    n_known_frames_for_test: null
-    path_manager_factory_PathManagerFactory_args:
-      silence_logs: true
   dataset_map_provider_JsonIndexDatasetMapProvider_args:
     category: ???
     task_str: singlesequence
@@ -91,14 +84,6 @@ data_source_ImplicitronDataSource_args:
       sort_frames: false
     path_manager_factory_PathManagerFactory_args:
       silence_logs: true
-  dataset_map_provider_LlffDatasetMapProvider_args:
-    base_dir: ???
-    object_name: ???
-    path_manager_factory_class_type: PathManagerFactory
-    n_known_frames_for_test: null
-    path_manager_factory_PathManagerFactory_args:
-      silence_logs: true
-    downscale_factor: 4
   dataset_map_provider_RenderedMeshDatasetMapProvider_args:
     num_views: 40
     data_file: null

projects/implicitron_trainer/tests/test_experiment.py

@@ -12,7 +12,6 @@ import unittest
 from pathlib import Path
 
 import torch
-
 from hydra import compose, initialize_config_dir
 from omegaconf import OmegaConf
 from projects.implicitron_trainer.impl.optimizer_factory import (

projects/implicitron_trainer/visualize_reconstruction.py

@@ -135,6 +135,7 @@ def visualize_reconstruction(
                 "device": device,
                 **render_flyaround_kwargs,
             }
+            # pyrefly: ignore [bad-argument-type]
             render_flyaround(**render_kwargs)
 
 

pytorch3d/__init__.py

@@ -6,4 +6,4 @@
 
 # pyre-unsafe
 
-__version__ = "0.7.8"
+__version__ = "0.7.9"

pytorch3d/common/workaround/symeig3x3.py

@@ -82,12 +82,15 @@ class _SymEig3x3(nn.Module):
         q = inputs_trace / 3.0
 
         # Calculate squared sum of elements outside the main diagonal / 2
-        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
-        p1 = ((inputs**2).sum(dim=(-1, -2)) - (inputs_diag**2).sum(-1)) / 2
-        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
-        p2 = ((inputs_diag - q[..., None]) ** 2).sum(dim=-1) + 2.0 * p1.clamp(self._eps)
+        p1 = (
+            torch.square(inputs).sum(dim=(-1, -2)) - torch.square(inputs_diag).sum(-1)
+        ) / 2
+        p2 = torch.square(inputs_diag - q[..., None]).sum(dim=-1) + 2.0 * p1.clamp(
+            self._eps
+        )
 
         p = torch.sqrt(p2 / 6.0)
+        # pyrefly: ignore [unsupported-operation]
         B = (inputs - q[..., None, None] * self._identity) / p[..., None, None]
 
         r = torch.det(B) / 2.0
@@ -104,7 +107,9 @@ class _SymEig3x3(nn.Module):
         # Soft dispatch between the degenerate case (diagonal A) and general.
         # diag_soft_cond -> 1.0 when p1 < 6 * eps and diag_soft_cond -> 0.0 otherwise.
         # We use 6 * eps to take into account the error accumulated during the p1 summation
-        diag_soft_cond = torch.exp(-((p1 / (6 * self._eps)) ** 2)).detach()[..., None]
+        diag_soft_cond = torch.exp(-torch.square(p1 / (6 * self._eps))).detach()[
+            ..., None
+        ]
 
         # Eigenvalues are the ordered elements of main diagonal in the degenerate case
         diag_eigenvals, _ = torch.sort(inputs_diag, dim=-1)
@@ -170,8 +175,10 @@ class _SymEig3x3(nn.Module):
         """
 
         # Find the eigenvector corresponding to alpha0, its eigenvalue is distinct
+        # pyrefly: ignore [unsupported-operation]
         ev0 = self._get_ev0(inputs - alpha0[..., None, None] * self._identity)
         u, v = self._get_uv(ev0)
+        # pyrefly: ignore [unsupported-operation]
         ev1 = self._get_ev1(inputs - alpha1[..., None, None] * self._identity, u, v)
         # Third eigenvector is computed as the cross-product of the other two
         ev2 = torch.cross(ev0, ev1, dim=-1)
@@ -199,8 +206,7 @@ class _SymEig3x3(nn.Module):
             cross_products[..., :1, :]
         )
 
-        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
-        norms_sq = (cross_products**2).sum(dim=-1)
+        norms_sq = torch.square(cross_products).sum(dim=-1)
         max_norms_index = norms_sq.argmax(dim=-1)
 
         # Pick only the cross-product with highest squared norm for each input
@@ -247,6 +253,7 @@ class _SymEig3x3(nn.Module):
         """
 
         min_idx = w.abs().argmin(dim=-1)
+        # pyrefly: ignore [bad-index]
         rotation_2d = self._rotations_3d[min_idx].to(w)
 
         u = F.normalize((rotation_2d @ w[..., None])[..., 0], dim=-1)

pytorch3d/csrc/ball_query/ball_query.cu

@@ -32,7 +32,9 @@ __global__ void BallQueryKernel(
     at::PackedTensorAccessor64<int64_t, 3, at::RestrictPtrTraits> idxs,
     at::PackedTensorAccessor64<scalar_t, 3, at::RestrictPtrTraits> dists,
     const int64_t K,
-    const float radius2) {
+    const float radius,
+    const float radius2,
+    const bool skip_points_outside_cube) {
   const int64_t N = p1.size(0);
   const int64_t chunks_per_cloud = (1 + (p1.size(1) - 1) / blockDim.x);
   const int64_t chunks_to_do = N * chunks_per_cloud;
@@ -51,7 +53,19 @@ __global__ void BallQueryKernel(
     // Iterate over points in p2 until desired count is reached or
     // all points have been considered
     for (int64_t j = 0, count = 0; j < lengths2[n] && count < K; ++j) {
-      // Calculate the distance between the points
+      if (skip_points_outside_cube) {
+        bool is_within_radius = true;
+        // Filter when any one coordinate is already outside the radius
+        for (int d = 0; is_within_radius && d < D; ++d) {
+          scalar_t abs_diff = fabs(p1[n][i][d] - p2[n][j][d]);
+          is_within_radius = (abs_diff <= radius);
+        }
+        if (!is_within_radius) {
+          continue;
+        }
+      }
+
+      // Else, calculate the distance between the points and compare
       scalar_t dist2 = 0.0;
       for (int d = 0; d < D; ++d) {
         scalar_t diff = p1[n][i][d] - p2[n][j][d];
@@ -77,7 +91,8 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
     const at::Tensor& lengths1, // (N,)
     const at::Tensor& lengths2, // (N,)
     int K,
-    float radius) {
+    float radius,
+    bool skip_points_outside_cube) {
   // Check inputs are on the same device
   at::TensorArg p1_t{p1, "p1", 1}, p2_t{p2, "p2", 2},
       lengths1_t{lengths1, "lengths1", 3}, lengths2_t{lengths2, "lengths2", 4};
@@ -120,7 +135,9 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
             idxs.packed_accessor64<int64_t, 3, at::RestrictPtrTraits>(),
             dists.packed_accessor64<float, 3, at::RestrictPtrTraits>(),
             K_64,
-            radius2);
+            radius,
+            radius2,
+            skip_points_outside_cube);
       }));
 
   AT_CUDA_CHECK(cudaGetLastError());

pytorch3d/csrc/ball_query/ball_query.h

@@ -25,6 +25,9 @@
 //      within the radius
 //    radius: the radius around each point within which the neighbors need to be
 //      located
+//    skip_points_outside_cube: If true, reduce multiplications of float values
+//      by not explicitly calculating distances to points that fall outside the
+//      D-cube with side length (2*radius) centered at each point in p1.
 //
 // Returns:
 //    p1_neighbor_idx: LongTensor of shape (N, P1, K), where
@@ -46,7 +49,8 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     const int K,
-    const float radius);
+    const float radius,
+    const bool skip_points_outside_cube);
 
 // CUDA implementation
 std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
@@ -55,7 +59,8 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     const int K,
-    const float radius);
+    const float radius,
+    const bool skip_points_outside_cube);
 
 // Implementation which is exposed
 // Note: the backward pass reuses the KNearestNeighborBackward kernel
@@ -65,7 +70,8 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     int K,
-    float radius) {
+    float radius,
+    bool skip_points_outside_cube) {
   if (p1.is_cuda() || p2.is_cuda()) {
 #ifdef WITH_CUDA
     CHECK_CUDA(p1);
@@ -76,16 +82,20 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
         lengths1.contiguous(),
         lengths2.contiguous(),
         K,
-        radius);
+        radius,
+        skip_points_outside_cube);
 #else
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(p1);
+  CHECK_CPU(p2);
   return BallQueryCpu(
       p1.contiguous(),
       p2.contiguous(),
       lengths1.contiguous(),
       lengths2.contiguous(),
       K,
-      radius);
+      radius,
+      skip_points_outside_cube);
 }

pytorch3d/csrc/ball_query/ball_query_cpu.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <math.h>
 #include <torch/extension.h>
 #include <tuple>
 
@@ -15,7 +16,8 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     int K,
-    float radius) {
+    float radius,
+    bool skip_points_outside_cube) {
   const int N = p1.size(0);
   const int P1 = p1.size(1);
   const int D = p1.size(2);
@@ -37,6 +39,16 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
     const int64_t length2 = lengths2_a[n];
     for (int64_t i = 0; i < length1; ++i) {
       for (int64_t j = 0, count = 0; j < length2 && count < K; ++j) {
+        if (skip_points_outside_cube) {
+          bool is_within_radius = true;
+          for (int d = 0; is_within_radius && d < D; ++d) {
+            float abs_diff = fabs(p1_a[n][i][d] - p2_a[n][j][d]);
+            is_within_radius = (abs_diff <= radius);
+          }
+          if (!is_within_radius) {
+            continue;
+          }
+        }
         float dist2 = 0;
         for (int d = 0; d < D; ++d) {
           float diff = p1_a[n][i][d] - p2_a[n][j][d];

pytorch3d/csrc/blending/sigmoid_alpha_blend.h

@@ -98,6 +98,11 @@ at::Tensor SigmoidAlphaBlendBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(distances);
+  CHECK_CPU(pix_to_face);
+  CHECK_CPU(alphas);
+  CHECK_CPU(grad_alphas);
+
   return SigmoidAlphaBlendBackwardCpu(
       grad_alphas, alphas, distances, pix_to_face, sigma);
 }

pytorch3d/csrc/compositing/alpha_composite.cu

@@ -33,11 +33,11 @@ __global__ void alphaCompositeCudaForwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Iterate over each feature in each pixel
   for (int pid = tid; pid < num_pixels; pid += num_threads) {
@@ -83,11 +83,11 @@ __global__ void alphaCompositeCudaBackwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Parallelize over each feature in each pixel in images of size H * W,
   // for each image in the batch of size batch_size

pytorch3d/csrc/compositing/alpha_composite.h

@@ -74,6 +74,9 @@ torch::Tensor alphaCompositeForward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
     return alphaCompositeCpuForward(features, alphas, points_idx);
   }
 }
@@ -101,6 +104,11 @@ std::tuple<torch::Tensor, torch::Tensor> alphaCompositeBackward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(grad_outputs);
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
+
     return alphaCompositeCpuBackward(
         grad_outputs, features, alphas, points_idx);
   }

pytorch3d/csrc/compositing/norm_weighted_sum.cu

@@ -33,11 +33,11 @@ __global__ void weightedSumNormCudaForwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Parallelize over each feature in each pixel in images of size H * W,
   // for each image in the batch of size batch_size
@@ -96,11 +96,11 @@ __global__ void weightedSumNormCudaBackwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * W * H;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Parallelize over each feature in each pixel in images of size H * W,
   // for each image in the batch of size batch_size

pytorch3d/csrc/compositing/norm_weighted_sum.h

@@ -73,6 +73,10 @@ torch::Tensor weightedSumNormForward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
+
     return weightedSumNormCpuForward(features, alphas, points_idx);
   }
 }
@@ -100,6 +104,11 @@ std::tuple<torch::Tensor, torch::Tensor> weightedSumNormBackward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(grad_outputs);
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
+
     return weightedSumNormCpuBackward(
         grad_outputs, features, alphas, points_idx);
   }

pytorch3d/csrc/compositing/weighted_sum.cu

@@ -31,11 +31,11 @@ __global__ void weightedSumCudaForwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Parallelize over each feature in each pixel in images of size H * W,
   // for each image in the batch of size batch_size
@@ -78,11 +78,11 @@ __global__ void weightedSumCudaBackwardKernel(
   const int64_t W = points_idx.size(3);
 
   // Get the batch and index
-  const int batch = blockIdx.x;
+  const auto batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const int num_threads = gridDim.y * blockDim.x;
-  const int tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.y * blockDim.x;
+  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
 
   // Iterate over each pixel to compute the contribution to the
   // gradient for the features and weights

pytorch3d/csrc/compositing/weighted_sum.h

@@ -72,6 +72,9 @@ torch::Tensor weightedSumForward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
     return weightedSumCpuForward(features, alphas, points_idx);
   }
 }
@@ -98,6 +101,11 @@ std::tuple<torch::Tensor, torch::Tensor> weightedSumBackward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(grad_outputs);
+    CHECK_CPU(features);
+    CHECK_CPU(alphas);
+    CHECK_CPU(points_idx);
+
     return weightedSumCpuBackward(grad_outputs, features, alphas, points_idx);
   }
 }

pytorch3d/csrc/ext.cpp

@@ -8,7 +8,6 @@
 
 // clang-format off
 #include "./pulsar/global.h" // Include before <torch/extension.h>.
-#include <torch/extension.h>
 // clang-format on
 #include "./pulsar/pytorch/renderer.h"
 #include "./pulsar/pytorch/tensor_util.h"
@@ -99,22 +98,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("marching_cubes", &MarchingCubes);
 
   // Pulsar.
-  // Pulsar not enabled on AMD.
 #ifdef PULSAR_LOGGING_ENABLED
   c10::ShowLogInfoToStderr();
 #endif
   py::class_<
       pulsar::pytorch::Renderer,
       std::shared_ptr<pulsar::pytorch::Renderer>>(m, "PulsarRenderer")
-      .def(py::init<
-           const uint&,
-           const uint&,
-           const uint&,
-           const bool&,
-           const bool&,
-           const float&,
-           const uint&,
-           const uint&>())
+      .def(
+          py::init<
+              const uint&,
+              const uint&,
+              const uint&,
+              const bool&,
+              const bool&,
+              const float&,
+              const uint&,
+              const uint&>())
       .def(
           "__eq__",
           [](const pulsar::pytorch::Renderer& a,

pytorch3d/csrc/face_areas_normals/face_areas_normals.h

@@ -60,6 +60,8 @@ std::tuple<at::Tensor, at::Tensor> FaceAreasNormalsForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(verts);
+  CHECK_CPU(faces);
   return FaceAreasNormalsForwardCpu(verts, faces);
 }
 
@@ -80,5 +82,9 @@ at::Tensor FaceAreasNormalsBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(grad_areas);
+  CHECK_CPU(grad_normals);
+  CHECK_CPU(verts);
+  CHECK_CPU(faces);
   return FaceAreasNormalsBackwardCpu(grad_areas, grad_normals, verts, faces);
 }

pytorch3d/csrc/gather_scatter/gather_scatter.cu

@@ -20,14 +20,14 @@ __global__ void GatherScatterCudaKernel(
     const size_t V,
     const size_t D,
     const size_t E) {
-  const int tid = threadIdx.x;
+  const auto tid = threadIdx.x;
 
   // Reverse the vertex order if backward.
   const int v0_idx = backward ? 1 : 0;
   const int v1_idx = backward ? 0 : 1;
 
   // Edges are split evenly across the blocks.
-  for (int e = blockIdx.x; e < E; e += gridDim.x) {
+  for (auto e = blockIdx.x; e < E; e += gridDim.x) {
     // Get indices of vertices which form the edge.
     const int64_t v0 = edges[2 * e + v0_idx];
     const int64_t v1 = edges[2 * e + v1_idx];
@@ -35,7 +35,7 @@ __global__ void GatherScatterCudaKernel(
     // Split vertex features evenly across threads.
     // This implementation will be quite wasteful when D<128 since there will be
     // a lot of threads doing nothing.
-    for (int d = tid; d < D; d += blockDim.x) {
+    for (auto d = tid; d < D; d += blockDim.x) {
       const float val = input[v1 * D + d];
       float* address = output + v0 * D + d;
       atomicAdd(address, val);

pytorch3d/csrc/gather_scatter/gather_scatter.h

@@ -53,5 +53,7 @@ at::Tensor GatherScatter(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(input);
+  CHECK_CPU(edges);
   return GatherScatterCpu(input, edges, directed, backward);
 }

pytorch3d/csrc/interp_face_attrs/interp_face_attrs.cu

@@ -20,8 +20,8 @@ __global__ void InterpFaceAttrsForwardKernel(
     const size_t P,
     const size_t F,
     const size_t D) {
-  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
-  const int num_threads = blockDim.x * gridDim.x;
+  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const auto num_threads = blockDim.x * gridDim.x;
   for (int pd = tid; pd < P * D; pd += num_threads) {
     const int p = pd / D;
     const int d = pd % D;
@@ -93,8 +93,8 @@ __global__ void InterpFaceAttrsBackwardKernel(
     const size_t P,
     const size_t F,
     const size_t D) {
-  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
-  const int num_threads = blockDim.x * gridDim.x;
+  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const auto num_threads = blockDim.x * gridDim.x;
   for (int pd = tid; pd < P * D; pd += num_threads) {
     const int p = pd / D;
     const int d = pd % D;

pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h

@@ -57,6 +57,8 @@ at::Tensor InterpFaceAttrsForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(face_attrs);
+  CHECK_CPU(barycentric_coords);
   return InterpFaceAttrsForwardCpu(pix_to_face, barycentric_coords, face_attrs);
 }
 
@@ -106,6 +108,9 @@ std::tuple<at::Tensor, at::Tensor> InterpFaceAttrsBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(face_attrs);
+  CHECK_CPU(barycentric_coords);
+  CHECK_CPU(grad_pix_attrs);
   return InterpFaceAttrsBackwardCpu(
       pix_to_face, barycentric_coords, face_attrs, grad_pix_attrs);
 }

pytorch3d/csrc/iou_box3d/iou_box3d.h

@@ -44,5 +44,7 @@ inline std::tuple<at::Tensor, at::Tensor> IoUBox3D(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(boxes1);
+  CHECK_CPU(boxes2);
   return IoUBox3DCpu(boxes1.contiguous(), boxes2.contiguous());
 }

pytorch3d/csrc/iou_box3d/iou_box3d_cpu.cpp

@@ -7,7 +7,6 @@
  */
 
 #include <torch/extension.h>
-#include <list>
 #include <tuple>
 #include "iou_box3d/iou_utils.h"
 

pytorch3d/csrc/knn/knn.h

@@ -74,6 +74,8 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborIdx(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(p1);
+  CHECK_CPU(p2);
   return KNearestNeighborIdxCpu(p1, p2, lengths1, lengths2, norm, K);
 }
 
@@ -140,6 +142,8 @@ std::tuple<at::Tensor, at::Tensor> KNearestNeighborBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(p1);
+  CHECK_CPU(p2);
   return KNearestNeighborBackwardCpu(
       p1, p2, lengths1, lengths2, idxs, norm, grad_dists);
 }

pytorch3d/csrc/marching_cubes/marching_cubes.h

@@ -58,5 +58,6 @@ inline std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubes(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(vol);
   return MarchingCubesCpu(vol.contiguous(), isolevel);
 }

pytorch3d/csrc/packed_to_padded_tensor/packed_to_padded_tensor.h

@@ -88,6 +88,8 @@ at::Tensor PackedToPadded(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(inputs_packed);
+  CHECK_CPU(first_idxs);
   return PackedToPaddedCpu(inputs_packed, first_idxs, max_size);
 }
 
@@ -105,5 +107,7 @@ at::Tensor PaddedToPacked(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(inputs_padded);
+  CHECK_CPU(first_idxs);
   return PaddedToPackedCpu(inputs_padded, first_idxs, num_inputs);
 }

pytorch3d/csrc/point_mesh/point_mesh_cpu.cpp

@@ -174,8 +174,8 @@ std::tuple<at::Tensor, at::Tensor> HullHullDistanceForwardCpu(
   at::Tensor idxs = at::zeros({A_N,}, as_first_idx.options());
   // clang-format on
 
-  auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
-  auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
+  auto as_a = as.accessor<float, H1 == 1 ? 2 : 3>();
+  auto bs_a = bs.accessor<float, H2 == 1 ? 2 : 3>();
   auto as_first_idx_a = as_first_idx.accessor<int64_t, 1>();
   auto bs_first_idx_a = bs_first_idx.accessor<int64_t, 1>();
   auto dists_a = dists.accessor<float, 1>();
@@ -230,10 +230,10 @@ std::tuple<at::Tensor, at::Tensor> HullHullDistanceBackwardCpu(
   at::Tensor grad_as = at::zeros_like(as);
   at::Tensor grad_bs = at::zeros_like(bs);
 
-  auto as_a = as.accessor < float, H1 == 1 ? 2 : 3 > ();
-  auto bs_a = bs.accessor < float, H2 == 1 ? 2 : 3 > ();
-  auto grad_as_a = grad_as.accessor < float, H1 == 1 ? 2 : 3 > ();
-  auto grad_bs_a = grad_bs.accessor < float, H2 == 1 ? 2 : 3 > ();
+  auto as_a = as.accessor<float, H1 == 1 ? 2 : 3>();
+  auto bs_a = bs.accessor<float, H2 == 1 ? 2 : 3>();
+  auto grad_as_a = grad_as.accessor<float, H1 == 1 ? 2 : 3>();
+  auto grad_bs_a = grad_bs.accessor<float, H2 == 1 ? 2 : 3>();
   auto idx_bs_a = idx_bs.accessor<int64_t, 1>();
   auto grad_dists_a = grad_dists.accessor<float, 1>();
 

pytorch3d/csrc/point_mesh/point_mesh_cuda.cu

@@ -110,7 +110,7 @@ __global__ void DistanceForwardKernel(
     __syncthreads();
 
     // Perform reduction in shared memory.
-    for (int s = blockDim.x / 2; s > 32; s >>= 1) {
+    for (auto s = blockDim.x / 2; s > 32; s >>= 1) {
       if (tid < s) {
         if (min_dists[tid] > min_dists[tid + s]) {
           min_dists[tid] = min_dists[tid + s];
@@ -502,8 +502,8 @@ __global__ void PointFaceArrayForwardKernel(
   const float3* tris_f3 = (float3*)tris;
 
   // Parallelize over P * S computations
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int t_i = tid; t_i < P * T; t_i += num_threads) {
     const int t = t_i / P; // segment index.
@@ -576,8 +576,8 @@ __global__ void PointFaceArrayBackwardKernel(
   const float3* tris_f3 = (float3*)tris;
 
   // Parallelize over P * S computations
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int t_i = tid; t_i < P * T; t_i += num_threads) {
     const int t = t_i / P; // triangle index.
@@ -683,8 +683,8 @@ __global__ void PointEdgeArrayForwardKernel(
   float3* segms_f3 = (float3*)segms;
 
   // Parallelize over P * S computations
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int t_i = tid; t_i < P * S; t_i += num_threads) {
     const int s = t_i / P; // segment index.
@@ -752,8 +752,8 @@ __global__ void PointEdgeArrayBackwardKernel(
   float3* segms_f3 = (float3*)segms;
 
   // Parallelize over P * S computations
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int t_i = tid; t_i < P * S; t_i += num_threads) {
     const int s = t_i / P; // segment index.

pytorch3d/csrc/point_mesh/point_mesh_cuda.h

@@ -88,6 +88,10 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(points_first_idx);
+  CHECK_CPU(tris);
+  CHECK_CPU(tris_first_idx);
   return PointFaceDistanceForwardCpu(
       points, points_first_idx, tris, tris_first_idx, min_triangle_area);
 }
@@ -143,6 +147,10 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(tris);
+  CHECK_CPU(idx_points);
+  CHECK_CPU(grad_dists);
   return PointFaceDistanceBackwardCpu(
       points, tris, idx_points, grad_dists, min_triangle_area);
 }
@@ -221,6 +229,10 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(points_first_idx);
+  CHECK_CPU(tris);
+  CHECK_CPU(tris_first_idx);
   return FacePointDistanceForwardCpu(
       points, points_first_idx, tris, tris_first_idx, min_triangle_area);
 }
@@ -277,6 +289,10 @@ std::tuple<torch::Tensor, torch::Tensor> FacePointDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(tris);
+  CHECK_CPU(idx_tris);
+  CHECK_CPU(grad_dists);
   return FacePointDistanceBackwardCpu(
       points, tris, idx_tris, grad_dists, min_triangle_area);
 }
@@ -346,6 +362,10 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(points_first_idx);
+  CHECK_CPU(segms);
+  CHECK_CPU(segms_first_idx);
   return PointEdgeDistanceForwardCpu(
       points, points_first_idx, segms, segms_first_idx, max_points);
 }
@@ -396,6 +416,10 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(segms);
+  CHECK_CPU(idx_points);
+  CHECK_CPU(grad_dists);
   return PointEdgeDistanceBackwardCpu(points, segms, idx_points, grad_dists);
 }
 
@@ -464,6 +488,10 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(points_first_idx);
+  CHECK_CPU(segms);
+  CHECK_CPU(segms_first_idx);
   return EdgePointDistanceForwardCpu(
       points, points_first_idx, segms, segms_first_idx, max_segms);
 }
@@ -514,6 +542,10 @@ std::tuple<torch::Tensor, torch::Tensor> EdgePointDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(segms);
+  CHECK_CPU(idx_segms);
+  CHECK_CPU(grad_dists);
   return EdgePointDistanceBackwardCpu(points, segms, idx_segms, grad_dists);
 }
 
@@ -567,6 +599,8 @@ torch::Tensor PointFaceArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(tris);
   return PointFaceArrayDistanceForwardCpu(points, tris, min_triangle_area);
 }
 
@@ -613,6 +647,9 @@ std::tuple<torch::Tensor, torch::Tensor> PointFaceArrayDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(tris);
+  CHECK_CPU(grad_dists);
   return PointFaceArrayDistanceBackwardCpu(
       points, tris, grad_dists, min_triangle_area);
 }
@@ -661,6 +698,8 @@ torch::Tensor PointEdgeArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(segms);
   return PointEdgeArrayDistanceForwardCpu(points, segms);
 }
 
@@ -703,5 +742,8 @@ std::tuple<torch::Tensor, torch::Tensor> PointEdgeArrayDistanceBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(segms);
+  CHECK_CPU(grad_dists);
   return PointEdgeArrayDistanceBackwardCpu(points, segms, grad_dists);
 }

pytorch3d/csrc/points_to_volumes/points_to_volumes.h

@@ -104,6 +104,12 @@ inline void PointsToVolumesForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points_3d);
+  CHECK_CPU(points_features);
+  CHECK_CPU(volume_densities);
+  CHECK_CPU(volume_features);
+  CHECK_CPU(grid_sizes);
+  CHECK_CPU(mask);
   PointsToVolumesForwardCpu(
       points_3d,
       points_features,
@@ -183,6 +189,14 @@ inline void PointsToVolumesBackward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points_3d);
+  CHECK_CPU(points_features);
+  CHECK_CPU(grid_sizes);
+  CHECK_CPU(mask);
+  CHECK_CPU(grad_volume_densities);
+  CHECK_CPU(grad_volume_features);
+  CHECK_CPU(grad_points_3d);
+  CHECK_CPU(grad_points_features);
   PointsToVolumesBackwardCpu(
       points_3d,
       points_features,

pytorch3d/csrc/pulsar/global.h

@@ -15,8 +15,8 @@
 #endif
 
 #if defined(_WIN64) || defined(_WIN32)
-#define uint unsigned int
-#define ushort unsigned short
+using uint = unsigned int;
+using ushort = unsigned short;
 #endif
 
 #include "./logging.h" // <- include before torch/extension.h

pytorch3d/csrc/pulsar/gpu/commands.h

@@ -125,29 +125,45 @@ INLINE DEVICE float3 WARP_SUM_FLOAT3(
 // Floating point.
 // #define FMUL(a, b) __fmul_rn((a), (b))
 #define FMUL(a, b) ((a) * (b))
-#define FDIV(a, b) __fdiv_rn((a), (b))
 // #define FSUB(a, b) __fsub_rn((a), (b))
 #define FSUB(a, b) ((a) - (b))
+// HIP has no per-instruction rounding-mode override (round-to-nearest-even is
+// the default at the hardware level), so the CUDA *_rn intrinsics have no
+// runtime equivalent and we use plain operators. The HIP/clang compiler may
+// fuse a+b*c into a single-rounding FMA where CUDA's _rn would have prevented
+// it; if that becomes a numerical issue, add `-ffp-contract=off` to the pulsar
+// nvcc_args in setup.py.
+#if defined(USE_ROCM)
+#define FDIV(a, b) ((a) / (b))
+#define FADD(a, b) ((a) + (b))
+#define FSQRT(a) sqrtf(a)
+#define FPOW(a, b) powf((a), (b))
+#define FSATURATE(x) fmaxf(0.0f, fminf(1.0f, (x)))
+#define FMA(x, y, z) fmaf((x), (y), (z))
+#define FRCP(x) (1.0f / (x))
+#else
+#define FDIV(a, b) __fdiv_rn((a), (b))
 #define FADD(a, b) __fadd_rn((a), (b))
 #define FSQRT(a) __fsqrt_rn(a)
+#define FPOW(a, b) __powf((a), (b))
+#define FSATURATE(x) __saturatef(x)
+/** Calculates x*y+z. */
+#define FMA(x, y, z) __fmaf_rn((x), (y), (z))
+#define FRCP(x) __frcp_rn(x)
+#endif
 #define FEXP(a) fasterexp(a)
 #define FLN(a) fasterlog(a)
-#define FPOW(a, b) __powf((a), (b))
 #define FMAX(a, b) fmax((a), (b))
 #define FMIN(a, b) fmin((a), (b))
 #define FCEIL(a) ceilf(a)
 #define FFLOOR(a) floorf(a)
 #define FROUND(x) nearbyintf(x)
-#define FSATURATE(x) __saturatef(x)
 #define FABS(a) abs(a)
 #define IASF(a, loc) (loc) = __int_as_float(a)
 #define FASI(a, loc) (loc) = __float_as_int(a)
 #define FABSLEQAS(a, b, c) \
   ((a) <= (b) ? FSUB((b), (a)) <= (c) : FSUB((a), (b)) < (c))
-/** Calculates x*y+z. */
-#define FMA(x, y, z) __fmaf_rn((x), (y), (z))
 #define I2F(a) __int2float_rn(a)
-#define FRCP(x) __frcp_rn(x)
 #if !defined(USE_ROCM)
 __device__ static float atomicMax(float* address, float val) {
   int* address_as_i = (int*)address;
@@ -201,8 +217,16 @@ __device__ static float atomicMin(float* address, float val) {
   ATOMICADD(&((PTR)->x), VAL.x); \
   ATOMICADD(&((PTR)->y), VAL.y); \
   ATOMICADD(&((PTR)->z), VAL.z);
-#if (CUDART_VERSION >= 10000) && (__CUDA_ARCH__ >= 600)
+#if !defined(USE_ROCM) && (CUDART_VERSION >= 10000) && (__CUDA_ARCH__ >= 600)
 #define ATOMICADD_B(PTR, VAL) atomicAdd_block((PTR), (VAL))
+#elif defined(USE_ROCM)
+// HIP has no atomicAdd_block, but the semantic equivalent is a relaxed
+// fetch_add scoped to the workgroup (HIP's name for a CUDA thread block).
+// This avoids device-wide L2-coherent atomics for what are block-local
+// counters in pulsar's inner sphere-loading loop.
+#define ATOMICADD_B(PTR, VAL) \
+  __hip_atomic_fetch_add(     \
+      (PTR), (VAL), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_WORKGROUP)
 #else
 #define ATOMICADD_B(PTR, VAL) ATOMICADD(PTR, VAL)
 #endif
@@ -417,7 +441,7 @@ __device__ static float atomicMin(float* address, float val) {
       (OUT_PTR),              \
       (NUM_SELECTED_PTR),     \
       (NUM_ITEMS),            \
-      stream = (STREAM));
+      (STREAM));
 
 #define COPY_HOST_DEV(PTR_D, PTR_H, TYPE, SIZE) \
   HANDLECUDA(cudaMemcpy(                        \

pytorch3d/csrc/pulsar/host/commands.h

@@ -357,11 +357,11 @@ void MAX_WS(
 //
 //
 #define END_PARALLEL() \
-  end_parallel :;      \
+  end_parallel:;       \
   }
 #define END_PARALLEL_NORET() }
 #define END_PARALLEL_2D() \
-  end_parallel :;         \
+  end_parallel:;          \
   }                       \
   }
 #define END_PARALLEL_2D_NORET() \

pytorch3d/csrc/pulsar/include/camera.h

@@ -70,11 +70,6 @@ struct CamGradInfo {
   float3 pixel_dir_y;
 };
 
-// TODO: remove once https://github.com/NVlabs/cub/issues/172 is resolved.
-struct IntWrapper {
-  int val;
-};
-
 } // namespace pulsar
 
 #endif

pytorch3d/csrc/pulsar/include/math.h

@@ -149,11 +149,6 @@ IHD CamGradInfo operator*(const CamGradInfo& a, const float& b) {
   return res;
 }
 
-IHD IntWrapper operator+(const IntWrapper& a, const IntWrapper& b) {
-  IntWrapper res;
-  res.val = a.val + b.val;
-  return res;
-}
 } // namespace pulsar
 
 #endif

pytorch3d/csrc/pulsar/include/renderer.backward.device.h

@@ -155,8 +155,8 @@ void backward(
         stream);
     CHECKLAUNCH();
     SUM_WS(
-        (IntWrapper*)(self->ids_sorted_d),
-        (IntWrapper*)(self->n_grad_contributions_d),
+        self->ids_sorted_d,
+        self->n_grad_contributions_d,
         static_cast<int>(num_balls),
         self->workspace_d,
         self->workspace_size,

pytorch3d/csrc/pulsar/include/renderer.construct.device.h

@@ -52,7 +52,7 @@ HOST void construct(
   self->cam.film_width = width;
   self->cam.film_height = height;
   self->max_num_balls = max_num_balls;
-  MALLOC(self->result_d, float, width* height* n_channels);
+  MALLOC(self->result_d, float, width * height * n_channels);
   self->cam.orthogonal_projection = orthogonal_projection;
   self->cam.right_handed = right_handed_system;
   self->cam.background_normalization_depth = background_normalization_depth;
@@ -93,7 +93,7 @@ HOST void construct(
   MALLOC(self->di_sorted_d, DrawInfo, max_num_balls);
   MALLOC(self->region_flags_d, char, max_num_balls);
   MALLOC(self->num_selected_d, size_t, 1);
-  MALLOC(self->forw_info_d, float, width* height * (3 + 2 * n_track));
+  MALLOC(self->forw_info_d, float, width * height * (3 + 2 * n_track));
   MALLOC(self->min_max_pixels_d, IntersectInfo, 1);
   MALLOC(self->grad_pos_d, float3, max_num_balls);
   MALLOC(self->grad_col_d, float, max_num_balls* n_channels);

pytorch3d/csrc/pulsar/include/renderer.destruct.device.h

@@ -18,68 +18,89 @@ namespace Renderer {
 
 template <bool DEV>
 HOST void destruct(Renderer* self) {
-  if (self->result_d != NULL)
+  if (self->result_d != NULL) {
     FREE(self->result_d);
+  }
   self->result_d = NULL;
-  if (self->min_depth_d != NULL)
+  if (self->min_depth_d != NULL) {
     FREE(self->min_depth_d);
+  }
   self->min_depth_d = NULL;
-  if (self->min_depth_sorted_d != NULL)
+  if (self->min_depth_sorted_d != NULL) {
     FREE(self->min_depth_sorted_d);
+  }
   self->min_depth_sorted_d = NULL;
-  if (self->ii_d != NULL)
+  if (self->ii_d != NULL) {
     FREE(self->ii_d);
+  }
   self->ii_d = NULL;
-  if (self->ii_sorted_d != NULL)
+  if (self->ii_sorted_d != NULL) {
     FREE(self->ii_sorted_d);
+  }
   self->ii_sorted_d = NULL;
-  if (self->ids_d != NULL)
+  if (self->ids_d != NULL) {
     FREE(self->ids_d);
+  }
   self->ids_d = NULL;
-  if (self->ids_sorted_d != NULL)
+  if (self->ids_sorted_d != NULL) {
     FREE(self->ids_sorted_d);
+  }
   self->ids_sorted_d = NULL;
-  if (self->workspace_d != NULL)
+  if (self->workspace_d != NULL) {
     FREE(self->workspace_d);
+  }
   self->workspace_d = NULL;
-  if (self->di_d != NULL)
+  if (self->di_d != NULL) {
     FREE(self->di_d);
+  }
   self->di_d = NULL;
-  if (self->di_sorted_d != NULL)
+  if (self->di_sorted_d != NULL) {
     FREE(self->di_sorted_d);
+  }
   self->di_sorted_d = NULL;
-  if (self->region_flags_d != NULL)
+  if (self->region_flags_d != NULL) {
     FREE(self->region_flags_d);
+  }
   self->region_flags_d = NULL;
-  if (self->num_selected_d != NULL)
+  if (self->num_selected_d != NULL) {
     FREE(self->num_selected_d);
+  }
   self->num_selected_d = NULL;
-  if (self->forw_info_d != NULL)
+  if (self->forw_info_d != NULL) {
     FREE(self->forw_info_d);
+  }
   self->forw_info_d = NULL;
-  if (self->min_max_pixels_d != NULL)
+  if (self->min_max_pixels_d != NULL) {
     FREE(self->min_max_pixels_d);
+  }
   self->min_max_pixels_d = NULL;
-  if (self->grad_pos_d != NULL)
+  if (self->grad_pos_d != NULL) {
     FREE(self->grad_pos_d);
+  }
   self->grad_pos_d = NULL;
-  if (self->grad_col_d != NULL)
+  if (self->grad_col_d != NULL) {
     FREE(self->grad_col_d);
+  }
   self->grad_col_d = NULL;
-  if (self->grad_rad_d != NULL)
+  if (self->grad_rad_d != NULL) {
     FREE(self->grad_rad_d);
+  }
   self->grad_rad_d = NULL;
-  if (self->grad_cam_d != NULL)
+  if (self->grad_cam_d != NULL) {
     FREE(self->grad_cam_d);
+  }
   self->grad_cam_d = NULL;
-  if (self->grad_cam_buf_d != NULL)
+  if (self->grad_cam_buf_d != NULL) {
     FREE(self->grad_cam_buf_d);
+  }
   self->grad_cam_buf_d = NULL;
-  if (self->grad_opy_d != NULL)
+  if (self->grad_opy_d != NULL) {
     FREE(self->grad_opy_d);
+  }
   self->grad_opy_d = NULL;
-  if (self->n_grad_contributions_d != NULL)
+  if (self->n_grad_contributions_d != NULL) {
     FREE(self->n_grad_contributions_d);
+  }
   self->n_grad_contributions_d = NULL;
 }
 

pytorch3d/csrc/pulsar/include/renderer.h

@@ -255,7 +255,7 @@ GLOBAL void calc_signature(
  * for every iteration through the loading loop every thread could add a
  * 'hit' to the buffer.
  */
-#define RENDER_BUFFER_SIZE RENDER_BLOCK_SIZE* RENDER_BLOCK_SIZE * 2
+#define RENDER_BUFFER_SIZE RENDER_BLOCK_SIZE * RENDER_BLOCK_SIZE * 2
 /**
  * The threshold after which the spheres that are in the render buffer
  * are rendered and the buffer is flushed.

pytorch3d/csrc/pulsar/include/renderer.norm_sphere_gradients.device.h

@@ -64,8 +64,9 @@ GLOBAL void norm_sphere_gradients(Renderer renderer, const int num_balls) {
   // The sphere only contributes to the camera gradients if it is
   // large enough in screen space.
   if (renderer.ids_sorted_d[idx] > 0 && ii.max.x >= ii.min.x + 3 &&
-      ii.max.y >= ii.min.y + 3)
+      ii.max.y >= ii.min.y + 3) {
     renderer.ids_sorted_d[idx] = 1;
+  }
   END_PARALLEL_NORET();
 };
 

pytorch3d/csrc/pulsar/include/renderer.render.device.h

@@ -139,8 +139,9 @@ GLOBAL void render(
       coord_y < cam_norm.film_border_top + cam_norm.film_height) {
     // Initialize the result.
     if (mode == 0u) {
-      for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id)
+      for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id) {
         result[c_id] = bg_col[c_id];
+      }
     } else {
       result[0] = 0.f;
     }
@@ -190,20 +191,22 @@ GLOBAL void render(
             "render|found intersection with sphere %u.\n",
             sphere_id_l[write_idx]);
       }
-      if (ii.min.x == MAX_USHORT)
+      if (ii.min.x == MAX_USHORT) {
         // This is an invalid sphere (out of image). These spheres have
         // maximum depth. Since we ordered the spheres by earliest possible
         // intersection depth we re certain that there will no other sphere
         // that is relevant after this one.
         loading_done = true;
+      }
     }
     // Reset n_pixels_done.
     n_pixels_done = 0;
     thread_block.sync(); // Make sure n_loaded is updated.
     if (n_loaded > RENDER_BUFFER_LOAD_THRESH) {
       // The load buffer is full enough. Draw.
-      if (thread_block.thread_rank() == 0)
+      if (thread_block.thread_rank() == 0) {
         n_balls_loaded += n_loaded;
+      }
       max_closest_possible_intersection = 0.f;
       // This excludes threads outside of the image boundary. Also, it reduces
       // block artifacts.
@@ -290,8 +293,9 @@ GLOBAL void render(
       uint warp_done = thread_warp.ballot(done);
       int warp_done_bit_cnt = POPC(warp_done);
 #endif //__CUDACC__ && __HIP_PLATFORM_AMD__
-      if (thread_warp.thread_rank() == 0)
+      if (thread_warp.thread_rank() == 0) {
         ATOMICADD_B(&n_pixels_done, warp_done_bit_cnt);
+      }
       // This sync is necessary to keep n_loaded until all threads are done with
       // painting.
       thread_block.sync();
@@ -299,8 +303,9 @@ GLOBAL void render(
     }
     thread_block.sync();
   }
-  if (thread_block.thread_rank() == 0)
+  if (thread_block.thread_rank() == 0) {
     n_balls_loaded += n_loaded;
+  }
   PULSAR_LOG_DEV_PIX(
       PULSAR_LOG_RENDER_PIX,
       "render|loaded %d balls in total.\n",
@@ -386,8 +391,9 @@ GLOBAL void render(
             static_cast<float>(tracker.get_n_hits());
   } else {
     float sm_d_normfac = FRCP(FMAX(sm_d, FEPS));
-    for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id)
+    for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id) {
       result[c_id] *= sm_d_normfac;
+    }
     int write_loc = (coord_y - cam_norm.film_border_top) * cam_norm.film_width *
             (3 + 2 * n_track) +
         (coord_x - cam_norm.film_border_left) * (3 + 2 * n_track);

pytorch3d/csrc/pulsar/pytorch/renderer.cpp

@@ -7,13 +7,10 @@
  */
 
 #include "./renderer.h"
-#include "../include/commands.h"
 #include "./camera.h"
 #include "./util.h"
 
-#include <ATen/ATen.h>
 #ifdef WITH_CUDA
-#include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #endif
 
@@ -860,8 +857,9 @@ std::tuple<torch::Tensor, torch::Tensor> Renderer::forward(
             ? (cudaStream_t) nullptr
 #endif
             : (cudaStream_t) nullptr);
-    if (mode == 1)
+    if (mode == 1) {
       results[batch_i] = results[batch_i].slice(2, 0, 1, 1);
+    }
     forw_infos[batch_i] = from_blob(
         this->renderer_vec[batch_i].forw_info_d,
         {this->renderer_vec[0].cam.film_height,

pytorch3d/csrc/pulsar/pytorch/renderer.h

@@ -128,8 +128,9 @@ struct Renderer {
     stream << "pulsar::Renderer[";
     // Device info.
     stream << self.device_type;
-    if (self.device_index != -1)
+    if (self.device_index != -1) {
       stream << ", ID " << self.device_index;
+    }
     stream << "]";
     return stream;
   }

pytorch3d/csrc/pulsar/warnings.cpp

@@ -6,9 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include "./global.h"
-#include "./logging.h"
-
 /**
  * A compilation unit to provide warnings about the code and avoid
  * repeated messages.

pytorch3d/csrc/rasterize_coarse/bitmask.cuh

@@ -25,7 +25,7 @@ class BitMask {
 
   // Use all threads in the current block to clear all bits of this BitMask
   __device__ void block_clear() {
-    for (int i = threadIdx.x; i < H * W * D; i += blockDim.x) {
+    for (auto i = threadIdx.x; i < H * W * D; i += blockDim.x) {
       data[i] = 0;
     }
     __syncthreads();

pytorch3d/csrc/rasterize_coarse/rasterize_coarse.cu

@@ -23,8 +23,8 @@ __global__ void TriangleBoundingBoxKernel(
     const float blur_radius,
     float* bboxes, // (4, F)
     bool* skip_face) { // (F,)
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int num_threads = blockDim.x * gridDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = blockDim.x * gridDim.x;
   const float sqrt_radius = sqrt(blur_radius);
   for (int f = tid; f < F; f += num_threads) {
     const float v0x = face_verts[f * 9 + 0 * 3 + 0];
@@ -56,8 +56,8 @@ __global__ void PointBoundingBoxKernel(
     const int P,
     float* bboxes, // (4, P)
     bool* skip_points) {
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int num_threads = blockDim.x * gridDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = blockDim.x * gridDim.x;
   for (int p = tid; p < P; p += num_threads) {
     const float x = points[p * 3 + 0];
     const float y = points[p * 3 + 1];
@@ -113,7 +113,7 @@ __global__ void RasterizeCoarseCudaKernel(
   const int chunks_per_batch = 1 + (E - 1) / chunk_size;
   const int num_chunks = N * chunks_per_batch;
 
-  for (int chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) {
+  for (auto chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) {
     const int batch_idx = chunk / chunks_per_batch; // batch index
     const int chunk_idx = chunk % chunks_per_batch;
     const int elem_chunk_start_idx = chunk_idx * chunk_size;
@@ -123,7 +123,7 @@ __global__ void RasterizeCoarseCudaKernel(
     const int64_t elem_stop_idx = elem_start_idx + elems_per_batch[batch_idx];
 
     // Have each thread handle a different face within the chunk
-    for (int e = threadIdx.x; e < chunk_size; e += blockDim.x) {
+    for (auto e = threadIdx.x; e < chunk_size; e += blockDim.x) {
       const int e_idx = elem_chunk_start_idx + e;
 
       // Check that we are still within the same element of the batch
@@ -170,7 +170,7 @@ __global__ void RasterizeCoarseCudaKernel(
     // Now we have processed every elem in the current chunk. We need to
     // count the number of elems in each bin so we can write the indices
     // out to global memory. We have each thread handle a different bin.
-    for (int byx = threadIdx.x; byx < num_bins_y * num_bins_x;
+    for (auto byx = threadIdx.x; byx < num_bins_y * num_bins_x;
          byx += blockDim.x) {
       const int by = byx / num_bins_x;
       const int bx = byx % num_bins_x;

pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu

@@ -260,8 +260,8 @@ __global__ void RasterizeMeshesNaiveCudaKernel(
     float* pix_dists,
     float* bary) {
   // Simple version: One thread per output pixel
-  int num_threads = gridDim.x * blockDim.x;
-  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  auto num_threads = gridDim.x * blockDim.x;
+  auto tid = blockDim.x * blockIdx.x + threadIdx.x;
 
   for (int i = tid; i < N * H * W; i += num_threads) {
     // Convert linear index to 3D index
@@ -446,8 +446,8 @@ __global__ void RasterizeMeshesBackwardCudaKernel(
 
   // Parallelize over each pixel in images of
   // size H * W, for each image in the batch of size N.
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int t_i = tid; t_i < N * H * W; t_i += num_threads) {
     // Convert linear index to 3D index
@@ -650,8 +650,8 @@ __global__ void RasterizeMeshesFineCudaKernel(
 ) {
   // This can be more than H * W if H or W are not divisible by bin_size.
   int num_pixels = N * BH * BW * bin_size * bin_size;
-  int num_threads = gridDim.x * blockDim.x;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  auto num_threads = gridDim.x * blockDim.x;
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int pid = tid; pid < num_pixels; pid += num_threads) {
     // Convert linear index into bin and pixel indices. We make the within

pytorch3d/csrc/rasterize_meshes/rasterize_meshes.h

@@ -138,6 +138,9 @@ RasterizeMeshesNaive(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(face_verts);
+    CHECK_CPU(mesh_to_face_first_idx);
+    CHECK_CPU(num_faces_per_mesh);
     return RasterizeMeshesNaiveCpu(
         face_verts,
         mesh_to_face_first_idx,
@@ -232,6 +235,11 @@ torch::Tensor RasterizeMeshesBackward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(face_verts);
+    CHECK_CPU(pix_to_face);
+    CHECK_CPU(grad_zbuf);
+    CHECK_CPU(grad_bary);
+    CHECK_CPU(grad_dists);
     return RasterizeMeshesBackwardCpu(
         face_verts,
         pix_to_face,
@@ -306,6 +314,9 @@ torch::Tensor RasterizeMeshesCoarse(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(face_verts);
+    CHECK_CPU(mesh_to_face_first_idx);
+    CHECK_CPU(num_faces_per_mesh);
     return RasterizeMeshesCoarseCpu(
         face_verts,
         mesh_to_face_first_idx,
@@ -423,6 +434,8 @@ RasterizeMeshesFine(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(face_verts);
+    CHECK_CPU(bin_faces);
     AT_ERROR("NOT IMPLEMENTED");
   }
 }

pytorch3d/csrc/rasterize_meshes/rasterize_meshes_cpu.cpp

@@ -8,7 +8,6 @@
 
 #include <torch/extension.h>
 #include <algorithm>
-#include <list>
 #include <thread>
 #include <tuple>
 #include "ATen/core/TensorAccessor.h"
@@ -106,6 +105,8 @@ auto ComputeFaceAreas(const torch::Tensor& face_verts) {
   return face_areas;
 }
 
+namespace {
+
 // Helper function to use with std::find_if to find the index of any
 // values in the top k struct which match a given idx.
 struct IsNeighbor {
@@ -118,7 +119,6 @@ struct IsNeighbor {
   int neighbor_idx;
 };
 
-namespace {
 void RasterizeMeshesNaiveCpu_worker(
     const int start_yi,
     const int end_yi,

pytorch3d/csrc/rasterize_points/rasterize_points.cu

@@ -97,8 +97,8 @@ __global__ void RasterizePointsNaiveCudaKernel(
     float* zbuf, // (N, H, W, K)
     float* pix_dists) { // (N, H, W, K)
   // Simple version: One thread per output pixel
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockDim.x * blockIdx.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockDim.x * blockIdx.x + threadIdx.x;
   for (int i = tid; i < N * H * W; i += num_threads) {
     // Convert linear index to 3D index
     const int n = i / (H * W); // Batch index
@@ -237,8 +237,8 @@ __global__ void RasterizePointsFineCudaKernel(
     float* pix_dists) { // (N, H, W, K)
   // This can be more than H * W if H or W are not divisible by bin_size.
   const int num_pixels = N * BH * BW * bin_size * bin_size;
-  const int num_threads = gridDim.x * blockDim.x;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto num_threads = gridDim.x * blockDim.x;
+  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
 
   for (int pid = tid; pid < num_pixels; pid += num_threads) {
     // Convert linear index into bin and pixel indices. We make the within
@@ -376,8 +376,8 @@ __global__ void RasterizePointsBackwardCudaKernel(
     float* grad_points) { // (P, 3)
   // Parallelized over each of K points per pixel, for each pixel in images of
   // size H * W, for each image in the batch of size N.
-  int num_threads = gridDim.x * blockDim.x;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  auto num_threads = gridDim.x * blockDim.x;
+  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
   for (int i = tid; i < N * H * W * K; i += num_threads) {
     // const int n = i / (H * W * K); // batch index (not needed).
     const int yxk = i % (H * W * K);

pytorch3d/csrc/rasterize_points/rasterize_points.h

@@ -91,6 +91,10 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsNaive(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(points);
+    CHECK_CPU(cloud_to_packed_first_idx);
+    CHECK_CPU(num_points_per_cloud);
+    CHECK_CPU(radius);
     return RasterizePointsNaiveCpu(
         points,
         cloud_to_packed_first_idx,
@@ -166,6 +170,10 @@ torch::Tensor RasterizePointsCoarse(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(points);
+    CHECK_CPU(cloud_to_packed_first_idx);
+    CHECK_CPU(num_points_per_cloud);
+    CHECK_CPU(radius);
     return RasterizePointsCoarseCpu(
         points,
         cloud_to_packed_first_idx,
@@ -232,6 +240,8 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> RasterizePointsFine(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(points);
+    CHECK_CPU(bin_points);
     AT_ERROR("NOT IMPLEMENTED");
   }
 }
@@ -284,6 +294,10 @@ torch::Tensor RasterizePointsBackward(
     AT_ERROR("Not compiled with GPU support");
 #endif
   } else {
+    CHECK_CPU(points);
+    CHECK_CPU(idxs);
+    CHECK_CPU(grad_zbuf);
+    CHECK_CPU(grad_dists);
     return RasterizePointsBackwardCpu(points, idxs, grad_zbuf, grad_dists);
   }
 }

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.cu

@@ -107,7 +107,8 @@ at::Tensor FarthestPointSamplingCuda(
     const at::Tensor& points, // (N, P, 3)
     const at::Tensor& lengths, // (N,)
     const at::Tensor& K, // (N,)
-    const at::Tensor& start_idxs) {
+    const at::Tensor& start_idxs,
+    const int64_t max_K_known = -1) {
   // Check inputs are on the same device
   at::TensorArg p_t{points, "points", 1}, lengths_t{lengths, "lengths", 2},
       k_t{K, "K", 3}, start_idxs_t{start_idxs, "start_idxs", 4};
@@ -129,7 +130,12 @@ at::Tensor FarthestPointSamplingCuda(
 
   const int64_t N = points.size(0);
   const int64_t P = points.size(1);
-  const int64_t max_K = at::max(K).item<int64_t>();
+  int64_t max_K;
+  if (max_K_known > 0) {
+    max_K = max_K_known;
+  } else {
+    max_K = at::max(K).item<int64_t>();
+  }
 
   // Initialize the output tensor with the sampled indices
   auto idxs = at::full({N, max_K}, -1, lengths.options());

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.h

@@ -43,7 +43,8 @@ at::Tensor FarthestPointSamplingCuda(
     const at::Tensor& points,
     const at::Tensor& lengths,
     const at::Tensor& K,
-    const at::Tensor& start_idxs);
+    const at::Tensor& start_idxs,
+    const int64_t max_K_known = -1);
 
 at::Tensor FarthestPointSamplingCpu(
     const at::Tensor& points,
@@ -56,17 +57,23 @@ at::Tensor FarthestPointSampling(
     const at::Tensor& points,
     const at::Tensor& lengths,
     const at::Tensor& K,
-    const at::Tensor& start_idxs) {
+    const at::Tensor& start_idxs,
+    const int64_t max_K_known = -1) {
   if (points.is_cuda() || lengths.is_cuda() || K.is_cuda()) {
 #ifdef WITH_CUDA
     CHECK_CUDA(points);
     CHECK_CUDA(lengths);
     CHECK_CUDA(K);
     CHECK_CUDA(start_idxs);
-    return FarthestPointSamplingCuda(points, lengths, K, start_idxs);
+    return FarthestPointSamplingCuda(
+        points, lengths, K, start_idxs, max_K_known);
 #else
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(points);
+  CHECK_CPU(lengths);
+  CHECK_CPU(K);
+  CHECK_CPU(start_idxs);
   return FarthestPointSamplingCpu(points, lengths, K, start_idxs);
 }

pytorch3d/csrc/sample_pdf/sample_pdf.h

@@ -71,6 +71,8 @@ inline void SamplePdf(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
+  CHECK_CPU(weights);
+  CHECK_CPU(outputs);
   CHECK_CONTIGUOUS(outputs);
   SamplePdfCpu(bins, weights, outputs, eps);
 }

pytorch3d/csrc/utils/dispatch.cuh

@@ -99,8 +99,7 @@ namespace {
 // and increment it via template recursion until it is equal to the run-time
 // argument N.
 template <
-    template <typename, int64_t>
-    class Kernel,
+    template <typename, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -124,8 +123,7 @@ struct DispatchKernelHelper1D {
 // 1D dispatch: Specialization when curN == maxN
 // We need this base case to avoid infinite template recursion.
 template <
-    template <typename, int64_t>
-    class Kernel,
+    template <typename, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -145,8 +143,7 @@ struct DispatchKernelHelper1D<Kernel, T, minN, maxN, maxN, Args...> {
 // the run-time values of N and M, at which point we dispatch to the run
 // method of the kernel.
 template <
-    template <typename, int64_t, int64_t>
-    class Kernel,
+    template <typename, int64_t, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -203,8 +200,7 @@ struct DispatchKernelHelper2D {
 
 // 2D dispatch, specialization for curN == maxN
 template <
-    template <typename, int64_t, int64_t>
-    class Kernel,
+    template <typename, int64_t, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -243,8 +239,7 @@ struct DispatchKernelHelper2D<
 
 // 2D dispatch, specialization for curM == maxM
 template <
-    template <typename, int64_t, int64_t>
-    class Kernel,
+    template <typename, int64_t, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -283,8 +278,7 @@ struct DispatchKernelHelper2D<
 
 // 2D dispatch, specialization for curN == maxN, curM == maxM
 template <
-    template <typename, int64_t, int64_t>
-    class Kernel,
+    template <typename, int64_t, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -313,8 +307,7 @@ struct DispatchKernelHelper2D<
 
 // This is the function we expect users to call to dispatch to 1D functions
 template <
-    template <typename, int64_t>
-    class Kernel,
+    template <typename, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,
@@ -330,8 +323,7 @@ void DispatchKernel1D(const int64_t N, Args... args) {
 
 // This is the function we expect users to call to dispatch to 2D functions
 template <
-    template <typename, int64_t, int64_t>
-    class Kernel,
+    template <typename, int64_t, int64_t> class Kernel,
     typename T,
     int64_t minN,
     int64_t maxN,

pytorch3d/csrc/utils/pytorch3d_cutils.h

@@ -15,3 +15,7 @@
 #define CHECK_CONTIGUOUS_CUDA(x) \
   CHECK_CUDA(x);                 \
   CHECK_CONTIGUOUS(x)
+#define CHECK_CPU(x)                    \
+  TORCH_CHECK(                          \
+      x.device().type() == torch::kCPU, \
+      "Cannot use CPU implementation: " #x " not on CPU.")

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/csrc/utils/warp_reduce.cuh

@@ -10,10 +10,20 @@
 #include <math.h>
 #include <cstdio>
 
-// Helper functions WarpReduceMin and WarpReduceMax used in .cu files
-// Starting in Volta, instructions are no longer synchronous within a warp.
-// We need to call __syncwarp() to sync the 32 threads in the warp
-// instead of all the threads in the block.
+// Helper functions WarpReduceMin and WarpReduceMax used in .cu files.
+// Starting in Volta, instructions are no longer synchronous within a warp,
+// so on CUDA __syncwarp() is required between dependent shared-memory
+// accesses in the unrolled tail reduction.
+//
+// On AMD/HIP no __syncwarp() is needed here: all wavefront lanes execute
+// in lockstep (AMD has no equivalent of NVIDIA's Independent Thread
+// Scheduling), and the AMDGPU memory model guarantees that LDS operations
+// issued by the same wavefront are observed in program order without an
+// explicit s_waitcnt — see the LLVM AMDGPU backend memory-model rules
+// (https://llvm.org/docs/AMDGPUUsage.html) and the HIP hardware-
+// implementation docs. This holds for both wave32 (RDNA, gfx10xx/11xx/
+// 12xx) and wave64 (CDNA, gfx9xx), so the USE_ROCM skip is
+// architecture-independent.
 
 template <typename scalar_t>
 __device__ void
@@ -23,8 +33,6 @@ WarpReduceMin(scalar_t* min_dists, int64_t* min_idxs, const size_t tid) {
     min_idxs[tid] = min_idxs[tid + 32];
     min_dists[tid] = min_dists[tid + 32];
   }
-// AMD does not use explicit syncwarp and instead automatically inserts memory
-// fences during compilation.
 #if !defined(USE_ROCM)
   __syncwarp();
 #endif

pytorch3d/datasets/r2n2/r2n2.py

@@ -200,6 +200,7 @@ class R2N2(ShapeNetBase):  # pragma: no cover
             ) % (shapenet_dir, ", ".join(synset_not_present))
             warnings.warn(msg)
 
+    # pyrefly: ignore [bad-override]
     def __getitem__(self, model_idx, view_idxs: Optional[List[int]] = None) -> Dict:
         """
         Read a model by the given index.
@@ -370,6 +371,7 @@ class R2N2(ShapeNetBase):  # pragma: no cover
         T = RT[3, :3]
         return R, T
 
+    # pyrefly: ignore [bad-override]
     def render(
         self,
         model_ids: Optional[List[str]] = None,

pytorch3d/datasets/r2n2/utils.py

@@ -62,8 +62,10 @@ def collate_batched_R2N2(batch: List[Dict]):  # pragma: no cover
     # all models have the same number of views V, stack the batches of
     # views of each model into a new batch of shape (N, V, H, W, 3).
     # Otherwise leave it as a list.
+    # pyrefly: ignore [not-iterable]
     if "images" in collated_dict:
         try:
+            # pyrefly: ignore [unsupported-operation]
             collated_dict["images"] = torch.stack(collated_dict["images"])
         except RuntimeError:
             print(
@@ -75,10 +77,14 @@ def collate_batched_R2N2(batch: List[Dict]):  # pragma: no cover
     # matrices and that all models have the same number of views V, stack each
     # type of matrices into a new batch of shape (N, V, ...).
     # Otherwise leave them as lists.
+    # pyrefly: ignore [not-iterable]
     if all(x in collated_dict for x in ["R", "T", "K"]):
         try:
+            # pyrefly: ignore [unsupported-operation]
             collated_dict["R"] = torch.stack(collated_dict["R"])  # (N, V, 3, 3)
+            # pyrefly: ignore [unsupported-operation]
             collated_dict["T"] = torch.stack(collated_dict["T"])  # (N, V, 3)
+            # pyrefly: ignore [unsupported-operation]
             collated_dict["K"] = torch.stack(collated_dict["K"])  # (N, V, 4, 4)
         except RuntimeError:
             print(
@@ -89,8 +95,10 @@ def collate_batched_R2N2(batch: List[Dict]):  # pragma: no cover
     # If collate_batched_meshes receives voxels and all models have the same
     # number of views V, stack the batches of voxels into a new batch of shape
     # (N, V, S, S, S), where S is the voxel size.
+    # pyrefly: ignore [not-iterable]
     if "voxels" in collated_dict:
         try:
+            # pyrefly: ignore [unsupported-operation]
             collated_dict["voxels"] = torch.stack(collated_dict["voxels"])
         except RuntimeError:
             print(
@@ -458,6 +466,8 @@ class BlenderCamera(CamerasBase):  # pragma: no cover
 
     def get_projection_transform(self, **kwargs) -> Transform3d:
         transform = Transform3d(device=self.device)
+        # pyre-fixme[29]: `Union[(self: TensorBase, dim0: int, dim1: int) -> Tensor,
+        #  Tensor, Module]` is not a function.
         transform._matrix = self.K.transpose(1, 2).contiguous()
         return transform
 

pytorch3d/datasets/shapenet_base.py

@@ -52,6 +52,7 @@ class ShapeNetBase(torch.utils.data.Dataset):  # pragma: no cover
         """
         return len(self.model_ids)
 
+    # pyrefly: ignore [bad-override-param-name]
     def __getitem__(self, idx) -> Dict:
         """
         Read a model by the given index. Need to be implemented for every child class
@@ -147,12 +148,17 @@ class ShapeNetBase(torch.utils.data.Dataset):  # pragma: no cover
         idxs = self._handle_render_inputs(model_ids, categories, sample_nums, idxs)
         # Use the getitem method which loads mesh + texture
         models = [self[idx] for idx in idxs]
+        # pyrefly: ignore [unsupported-operation]
         meshes = collate_batched_meshes(models)["mesh"]
+        # pyrefly: ignore [missing-attribute]
         if meshes.textures is None:
+            # pyrefly: ignore [missing-attribute]
             meshes.textures = TexturesVertex(
+                # pyrefly: ignore [missing-attribute]
                 verts_features=torch.ones_like(meshes.verts_padded(), device=device)
             )
 
+        # pyrefly: ignore [missing-attribute]
         meshes = meshes.to(device)
         cameras = kwargs.get("cameras", FoVPerspectiveCameras()).to(device)
         if len(cameras) != 1 and len(cameras) % len(meshes) != 0:

pytorch3d/datasets/utils.py

@@ -34,12 +34,14 @@ def collate_batched_meshes(batch: List[Dict]):  # pragma: no cover
     for k in batch[0].keys():
         collated_dict[k] = [d[k] for d in batch]
 
+    # pyrefly: ignore [unsupported-operation]
     collated_dict["mesh"] = None
     if {"verts", "faces"}.issubset(collated_dict.keys()):
         textures = None
         if "textures" in collated_dict:
             textures = TexturesAtlas(atlas=collated_dict["textures"])
 
+        # pyrefly: ignore [unsupported-operation]
         collated_dict["mesh"] = Meshes(
             verts=collated_dict["verts"],
             faces=collated_dict["faces"],

pytorch3d/implicitron/dataset/blender_dataset_map_provider.py

@@ -1,55 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# pyre-unsafe
-
-
-import torch
-from pytorch3d.implicitron.tools.config import registry
-
-from .load_blender import load_blender_data
-from .single_sequence_dataset import (
-    _interpret_blender_cameras,
-    SingleSceneDatasetMapProviderBase,
-)
-
-
-@registry.register
-class BlenderDatasetMapProvider(SingleSceneDatasetMapProviderBase):
-    """
-    Provides data for one scene from Blender synthetic dataset.
-    Uses the code in load_blender.py
-
-    Members:
-        base_dir: directory holding the data for the scene.
-        object_name: The name of the scene (e.g. "lego"). This is just used as a label.
-            It will typically be equal to the name of the directory self.base_dir.
-        path_manager_factory: Creates path manager which may be used for
-            interpreting paths.
-        n_known_frames_for_test: If set, training frames are included in the val
-            and test datasets, and this many random training frames are added to
-            each test batch. If not set, test batches each contain just a single
-            testing frame.
-    """
-
-    def _load_data(self) -> None:
-        path_manager = self.path_manager_factory.get()
-        images, poses, _, hwf, i_split = load_blender_data(
-            self.base_dir,
-            testskip=1,
-            path_manager=path_manager,
-        )
-        H, W, focal = hwf
-        images_masks = torch.from_numpy(images).permute(0, 3, 1, 2)
-
-        # pyre-ignore[16]
-        self.poses = _interpret_blender_cameras(poses, focal)
-        # pyre-ignore[16]
-        self.images = images_masks[:, :3]
-        # pyre-ignore[16]
-        self.fg_probabilities = images_masks[:, 3:4]
-        # pyre-ignore[16]
-        self.i_split = i_split

pytorch3d/implicitron/dataset/data_source.py

@@ -64,16 +64,12 @@ class ImplicitronDataSource(DataSourceBase):
     def pre_expand(cls) -> None:
         # use try/finally to bypass cinder's lazy imports
         try:
-            from .blender_dataset_map_provider import (  # noqa: F401
-                BlenderDatasetMapProvider,
-            )
             from .json_index_dataset_map_provider import (  # noqa: F401
                 JsonIndexDatasetMapProvider,
             )
             from .json_index_dataset_map_provider_v2 import (  # noqa: F401
                 JsonIndexDatasetMapProviderV2,
             )
-            from .llff_dataset_map_provider import LlffDatasetMapProvider  # noqa: F401
             from .rendered_mesh_dataset_map_provider import (  # noqa: F401
                 RenderedMeshDatasetMapProvider,
             )

pytorch3d/implicitron/dataset/dataset_base.py

@@ -21,7 +21,6 @@ from typing import (
 )
 
 import torch
-
 from pytorch3d.implicitron.dataset.frame_data import FrameData
 from pytorch3d.implicitron.dataset.utils import GenericWorkaround
 
@@ -102,6 +101,7 @@ class DatasetBase(GenericWorkaround, torch.utils.data.Dataset[FrameData]):
             # crashes without overriding __getitem__
             sequence_category = self[first_frame_idx].sequence_category
             c2seq[sequence_category].append(sequence_name)
+        # pyrefly: ignore [bad-return]
         return dict(c2seq)
 
     def sequence_frames_in_order(

pytorch3d/implicitron/dataset/frame_data.py

@@ -25,7 +25,6 @@ from typing import (
 
 import numpy as np
 import torch
-
 from pytorch3d.implicitron.dataset import orm_types, types
 from pytorch3d.implicitron.dataset.utils import (
     adjust_camera_to_bbox_crop_,
@@ -298,7 +297,11 @@ class FrameData(Mapping[str, Any]):
         depth_map = self.depth_map
         if depth_map is not None:
             clamp_bbox_xyxy_depth = rescale_bbox(
-                clamp_bbox_xyxy, tuple(depth_map.shape[-2:]), effective_image_size_hw
+                # pyrefly: ignore [bad-argument-type]
+                clamp_bbox_xyxy,
+                # pyrefly: ignore [bad-argument-type]
+                tuple(depth_map.shape[-2:]),
+                effective_image_size_hw,
             ).long()
             self.depth_map = crop_around_box(
                 depth_map,
@@ -309,7 +312,11 @@ class FrameData(Mapping[str, Any]):
         depth_mask = self.depth_mask
         if depth_mask is not None:
             clamp_bbox_xyxy_depth = rescale_bbox(
-                clamp_bbox_xyxy, tuple(depth_mask.shape[-2:]), effective_image_size_hw
+                # pyrefly: ignore [bad-argument-type]
+                clamp_bbox_xyxy,
+                # pyrefly: ignore [bad-argument-type]
+                tuple(depth_mask.shape[-2:]),
+                effective_image_size_hw,
             ).long()
             self.depth_mask = crop_around_box(
                 depth_mask,
@@ -454,6 +461,7 @@ class FrameDataBuilderBase(ReplaceableBase, Generic[FrameDataSubtype], ABC):
     """
 
     # To be initialised to FrameDataSubtype
+    # pyrefly: ignore [invalid-annotation]
     frame_data_type: ClassVar[Type[FrameDataSubtype]]
 
     @abstractmethod

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -38,7 +38,6 @@ from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
 from pytorch3d.renderer.camera_utils import join_cameras_as_batch
 from pytorch3d.renderer.cameras import CamerasBase
-
 from tqdm import tqdm
 
 
@@ -327,9 +326,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
                 assert os.path.normpath(
                     # pyre-ignore[16]
                     self.frame_annots[idx]["frame_annotation"].image.path
-                ) == os.path.normpath(
-                    path
-                ), f"Inconsistent frame indices {seq_name, frame_no, path}."
+                ) == os.path.normpath(path), (
+                    f"Inconsistent frame indices {seq_name, frame_no, path}."
+                )
             return idx
 
         dataset_idx = [
@@ -426,6 +425,7 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
                 raise ValueError("subsets not loaded")
             if is_known_frame_scalar(frame_type):
                 cameras.append(self[frame_idx].camera)
+        # pyrefly: ignore [bad-argument-type]
         return join_cameras_as_batch(cameras)
 
     def __getitem__(self, index) -> FrameData:

pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py

@@ -21,7 +21,6 @@ from pytorch3d.renderer.cameras import CamerasBase
 
 from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, PathManagerFactory
 from .json_index_dataset import JsonIndexDataset
-
 from .utils import (
     DATASET_TYPE_KNOWN,
     DATASET_TYPE_TEST,

pytorch3d/implicitron/dataset/json_index_dataset_map_provider_v2.py

@@ -18,7 +18,6 @@ from typing import Dict, List, Optional, Tuple, Type, Union
 
 import numpy as np
 from iopath.common.file_io import PathManager
-
 from omegaconf import DictConfig
 from pytorch3d.implicitron.dataset.dataset_map_provider import (
     DatasetMap,
@@ -31,7 +30,6 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.renderer.cameras import CamerasBase
 from tqdm import tqdm
 
@@ -313,6 +311,7 @@ class JsonIndexDatasetMapProviderV2(DatasetMapProviderBase):
                     subset_mapping["test"],
                 ) = self._extend_test_data_with_known_views(
                     subset_mapping,
+                    # pyrefly: ignore [bad-argument-type]
                     eval_batch_index,
                 )
 
@@ -324,6 +323,7 @@ class JsonIndexDatasetMapProviderV2(DatasetMapProviderBase):
                 try:
                     test_dataset.eval_batches = (
                         test_dataset.seq_frame_index_to_dataset_index(
+                            # pyrefly: ignore [bad-argument-type]
                             eval_batch_index,
                         )
                     )
@@ -337,6 +337,7 @@ class JsonIndexDatasetMapProviderV2(DatasetMapProviderBase):
                     )
                     test_dataset.eval_batches = (
                         test_dataset.seq_frame_index_to_dataset_index(
+                            # pyrefly: ignore [bad-argument-type]
                             eval_batch_index,
                             allow_missing_indices=True,
                             remove_missing_indices=True,

pytorch3d/implicitron/dataset/llff_dataset_map_provider.py

@@ -1,69 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
-
-# pyre-unsafe
-
-
-import numpy as np
-import torch
-from pytorch3d.implicitron.tools.config import registry
-
-from .load_llff import load_llff_data
-
-from .single_sequence_dataset import (
-    _interpret_blender_cameras,
-    SingleSceneDatasetMapProviderBase,
-)
-
-
-@registry.register
-class LlffDatasetMapProvider(SingleSceneDatasetMapProviderBase):
-    """
-    Provides data for one scene from the LLFF dataset.
-
-    Members:
-        base_dir: directory holding the data for the scene.
-        object_name: The name of the scene (e.g. "fern"). This is just used as a label.
-            It will typically be equal to the name of the directory self.base_dir.
-        path_manager_factory: Creates path manager which may be used for
-            interpreting paths.
-        n_known_frames_for_test: If set, training frames are included in the val
-            and test datasets, and this many random training frames are added to
-            each test batch. If not set, test batches each contain just a single
-            testing frame.
-        downscale_factor: determines image sizes.
-    """
-
-    downscale_factor: int = 4
-
-    def _load_data(self) -> None:
-        path_manager = self.path_manager_factory.get()
-        images, poses, _ = load_llff_data(
-            self.base_dir, factor=self.downscale_factor, path_manager=path_manager
-        )
-        hwf = poses[0, :3, -1]
-        poses = poses[:, :3, :4]
-
-        llffhold = 8
-        i_test = np.arange(images.shape[0])[::llffhold]
-        i_test_index = set(i_test.tolist())
-        i_train = np.array(
-            [i for i in np.arange(images.shape[0]) if i not in i_test_index]
-        )
-        i_split = (i_train, i_test, i_test)
-        H, W, focal = hwf
-        focal_ndc = 2 * focal / min(H, W)
-        images = torch.from_numpy(images).permute(0, 3, 1, 2)
-        poses = torch.from_numpy(poses)
-
-        # pyre-ignore[16]
-        self.poses = _interpret_blender_cameras(poses, focal_ndc)
-        # pyre-ignore[16]
-        self.images = images
-        # pyre-ignore[16]
-        self.fg_probabilities = None
-        # pyre-ignore[16]
-        self.i_split = i_split

pytorch3d/implicitron/dataset/load_blender.py

@@ -1,143 +0,0 @@
-# @lint-ignore-every LICENSELINT
-# Adapted from https://github.com/bmild/nerf/blob/master/load_blender.py
-# Copyright (c) 2020 bmild
-
-# pyre-unsafe
-import json
-import os
-
-import numpy as np
-import torch
-from PIL import Image
-
-
-def translate_by_t_along_z(t):
-    tform = np.eye(4).astype(np.float32)
-    tform[2][3] = t
-    return tform
-
-
-def rotate_by_phi_along_x(phi):
-    tform = np.eye(4).astype(np.float32)
-    tform[1, 1] = tform[2, 2] = np.cos(phi)
-    tform[1, 2] = -np.sin(phi)
-    tform[2, 1] = -tform[1, 2]
-    return tform
-
-
-def rotate_by_theta_along_y(theta):
-    tform = np.eye(4).astype(np.float32)
-    tform[0, 0] = tform[2, 2] = np.cos(theta)
-    tform[0, 2] = -np.sin(theta)
-    tform[2, 0] = -tform[0, 2]
-    return tform
-
-
-def pose_spherical(theta, phi, radius):
-    c2w = translate_by_t_along_z(radius)
-    c2w = rotate_by_phi_along_x(phi / 180.0 * np.pi) @ c2w
-    c2w = rotate_by_theta_along_y(theta / 180 * np.pi) @ c2w
-    c2w = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) @ c2w
-    return c2w
-
-
-def _local_path(path_manager, path):
-    if path_manager is None:
-        return path
-    return path_manager.get_local_path(path)
-
-
-def load_blender_data(
-    basedir,
-    half_res=False,
-    testskip=1,
-    debug=False,
-    path_manager=None,
-    focal_length_in_screen_space=False,
-):
-    splits = ["train", "val", "test"]
-    metas = {}
-    for s in splits:
-        path = os.path.join(basedir, f"transforms_{s}.json")
-        with open(_local_path(path_manager, path)) as fp:
-            metas[s] = json.load(fp)
-
-    all_imgs = []
-    all_poses = []
-    counts = [0]
-    for s in splits:
-        meta = metas[s]
-        imgs = []
-        poses = []
-        if s == "train" or testskip == 0:
-            skip = 1
-        else:
-            skip = testskip
-
-        for frame in meta["frames"][::skip]:
-            fname = os.path.join(basedir, frame["file_path"] + ".png")
-            imgs.append(np.array(Image.open(_local_path(path_manager, fname))))
-            poses.append(np.array(frame["transform_matrix"]))
-        imgs = (np.array(imgs) / 255.0).astype(np.float32)
-        poses = np.array(poses).astype(np.float32)
-        counts.append(counts[-1] + imgs.shape[0])
-        all_imgs.append(imgs)
-        all_poses.append(poses)
-
-    i_split = [np.arange(counts[i], counts[i + 1]) for i in range(3)]
-
-    imgs = np.concatenate(all_imgs, 0)
-    poses = np.concatenate(all_poses, 0)
-
-    H, W = imgs[0].shape[:2]
-    camera_angle_x = float(meta["camera_angle_x"])
-    if focal_length_in_screen_space:
-        focal = 0.5 * W / np.tan(0.5 * camera_angle_x)
-    else:
-        focal = 1 / np.tan(0.5 * camera_angle_x)
-
-    render_poses = torch.stack(
-        [
-            torch.from_numpy(pose_spherical(angle, -30.0, 4.0))
-            for angle in np.linspace(-180, 180, 40 + 1)[:-1]
-        ],
-        0,
-    )
-
-    # In debug mode, return extremely tiny images
-    if debug:
-        import cv2
-
-        H = H // 32
-        W = W // 32
-        if focal_length_in_screen_space:
-            focal = focal / 32.0
-        imgs = [
-            torch.from_numpy(
-                cv2.resize(imgs[i], dsize=(25, 25), interpolation=cv2.INTER_AREA)
-            )
-            for i in range(imgs.shape[0])
-        ]
-        imgs = torch.stack(imgs, 0)
-        poses = torch.from_numpy(poses)
-        return imgs, poses, render_poses, [H, W, focal], i_split
-
-    if half_res:
-        import cv2
-
-        # TODO: resize images using INTER_AREA (cv2)
-        H = H // 2
-        W = W // 2
-        if focal_length_in_screen_space:
-            focal = focal / 2.0
-        imgs = [
-            torch.from_numpy(
-                cv2.resize(imgs[i], dsize=(400, 400), interpolation=cv2.INTER_AREA)
-            )
-            for i in range(imgs.shape[0])
-        ]
-        imgs = torch.stack(imgs, 0)
-
-    poses = torch.from_numpy(poses)
-
-    return imgs, poses, render_poses, [H, W, focal], i_split

pytorch3d/implicitron/dataset/load_llff.py

@@ -1,336 +0,0 @@
-# @lint-ignore-every LICENSELINT
-# Adapted from https://github.com/bmild/nerf/blob/master/load_llff.py
-# Copyright (c) 2020 bmild
-
-# pyre-unsafe
-import logging
-import os
-import warnings
-
-import numpy as np
-
-from PIL import Image
-
-
-# Slightly modified version of LLFF data loading code
-#  see https://github.com/Fyusion/LLFF for original
-
-logger = logging.getLogger(__name__)
-
-
-def _minify(basedir, path_manager, factors=(), resolutions=()):
-    needtoload = False
-    for r in factors:
-        imgdir = os.path.join(basedir, "images_{}".format(r))
-        if not _exists(path_manager, imgdir):
-            needtoload = True
-    for r in resolutions:
-        imgdir = os.path.join(basedir, "images_{}x{}".format(r[1], r[0]))
-        if not _exists(path_manager, imgdir):
-            needtoload = True
-    if not needtoload:
-        return
-    assert path_manager is None
-
-    from subprocess import check_output
-
-    imgdir = os.path.join(basedir, "images")
-    imgs = [os.path.join(imgdir, f) for f in sorted(_ls(path_manager, imgdir))]
-    imgs = [f for f in imgs if f.endswith("JPG", "jpg", "png", "jpeg", "PNG")]
-    imgdir_orig = imgdir
-
-    wd = os.getcwd()
-
-    for r in factors + resolutions:
-        if isinstance(r, int):
-            name = "images_{}".format(r)
-            resizearg = "{}%".format(100.0 / r)
-        else:
-            name = "images_{}x{}".format(r[1], r[0])
-            resizearg = "{}x{}".format(r[1], r[0])
-        imgdir = os.path.join(basedir, name)
-        if os.path.exists(imgdir):
-            continue
-
-        logger.info(f"Minifying {r}, {basedir}")
-
-        os.makedirs(imgdir)
-        check_output("cp {}/* {}".format(imgdir_orig, imgdir), shell=True)
-
-        ext = imgs[0].split(".")[-1]
-        args = " ".join(
-            ["mogrify", "-resize", resizearg, "-format", "png", "*.{}".format(ext)]
-        )
-        logger.info(args)
-        os.chdir(imgdir)
-        check_output(args, shell=True)
-        os.chdir(wd)
-
-        if ext != "png":
-            check_output("rm {}/*.{}".format(imgdir, ext), shell=True)
-            logger.info("Removed duplicates")
-        logger.info("Done")
-
-
-def _load_data(
-    basedir, factor=None, width=None, height=None, load_imgs=True, path_manager=None
-):
-    poses_arr = np.load(
-        _local_path(path_manager, os.path.join(basedir, "poses_bounds.npy"))
-    )
-    poses = poses_arr[:, :-2].reshape([-1, 3, 5]).transpose([1, 2, 0])
-    bds = poses_arr[:, -2:].transpose([1, 0])
-
-    img0 = [
-        os.path.join(basedir, "images", f)
-        for f in sorted(_ls(path_manager, os.path.join(basedir, "images")))
-        if f.endswith("JPG") or f.endswith("jpg") or f.endswith("png")
-    ][0]
-
-    def imread(f):
-        return np.array(Image.open(f))
-
-    sh = imread(_local_path(path_manager, img0)).shape
-
-    sfx = ""
-
-    if factor is not None:
-        sfx = "_{}".format(factor)
-        _minify(basedir, path_manager, factors=[factor])
-        factor = factor
-    elif height is not None:
-        factor = sh[0] / float(height)
-        width = int(sh[1] / factor)
-        _minify(basedir, path_manager, resolutions=[[height, width]])
-        sfx = "_{}x{}".format(width, height)
-    elif width is not None:
-        factor = sh[1] / float(width)
-        height = int(sh[0] / factor)
-        _minify(basedir, path_manager, resolutions=[[height, width]])
-        sfx = "_{}x{}".format(width, height)
-    else:
-        factor = 1
-
-    imgdir = os.path.join(basedir, "images" + sfx)
-    if not _exists(path_manager, imgdir):
-        raise ValueError(f"{imgdir} does not exist, returning")
-
-    imgfiles = [
-        _local_path(path_manager, os.path.join(imgdir, f))
-        for f in sorted(_ls(path_manager, imgdir))
-        if f.endswith("JPG") or f.endswith("jpg") or f.endswith("png")
-    ]
-    if poses.shape[-1] != len(imgfiles):
-        raise ValueError(
-            "Mismatch between imgs {} and poses {} !!!!".format(
-                len(imgfiles), poses.shape[-1]
-            )
-        )
-
-    sh = imread(imgfiles[0]).shape
-    poses[:2, 4, :] = np.array(sh[:2]).reshape([2, 1])
-    poses[2, 4, :] = poses[2, 4, :] * 1.0 / factor
-
-    if not load_imgs:
-        return poses, bds
-
-    imgs = imgs = [imread(f)[..., :3] / 255.0 for f in imgfiles]
-    imgs = np.stack(imgs, -1)
-
-    logger.info(f"Loaded image data, shape {imgs.shape}")
-    return poses, bds, imgs
-
-
-def normalize(x):
-    denom = np.linalg.norm(x)
-    if denom < 0.001:
-        warnings.warn("unsafe normalize()")
-    return x / denom
-
-
-def viewmatrix(z, up, pos):
-    vec2 = normalize(z)
-    vec1_avg = up
-    vec0 = normalize(np.cross(vec1_avg, vec2))
-    vec1 = normalize(np.cross(vec2, vec0))
-    m = np.stack([vec0, vec1, vec2, pos], 1)
-    return m
-
-
-def ptstocam(pts, c2w):
-    tt = np.matmul(c2w[:3, :3].T, (pts - c2w[:3, 3])[..., np.newaxis])[..., 0]
-    return tt
-
-
-def poses_avg(poses):
-    hwf = poses[0, :3, -1:]
-
-    center = poses[:, :3, 3].mean(0)
-    vec2 = normalize(poses[:, :3, 2].sum(0))
-    up = poses[:, :3, 1].sum(0)
-    c2w = np.concatenate([viewmatrix(vec2, up, center), hwf], 1)
-
-    return c2w
-
-
-def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, rots, N):
-    render_poses = []
-    rads = np.array(list(rads) + [1.0])
-    hwf = c2w[:, 4:5]
-
-    for theta in np.linspace(0.0, 2.0 * np.pi * rots, N + 1)[:-1]:
-        c = np.dot(
-            c2w[:3, :4],
-            np.array([np.cos(theta), -np.sin(theta), -np.sin(theta * zrate), 1.0])
-            * rads,
-        )
-        z = normalize(c - np.dot(c2w[:3, :4], np.array([0, 0, -focal, 1.0])))
-        render_poses.append(np.concatenate([viewmatrix(z, up, c), hwf], 1))
-    return render_poses
-
-
-def recenter_poses(poses):
-    poses_ = poses + 0
-    bottom = np.reshape([0, 0, 0, 1.0], [1, 4])
-    c2w = poses_avg(poses)
-    c2w = np.concatenate([c2w[:3, :4], bottom], -2)
-    bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1])
-    poses = np.concatenate([poses[:, :3, :4], bottom], -2)
-
-    poses = np.linalg.inv(c2w) @ poses
-    poses_[:, :3, :4] = poses[:, :3, :4]
-    poses = poses_
-    return poses
-
-
-def spherify_poses(poses, bds):
-    def add_row_to_homogenize_transform(p):
-        r"""Add the last row to homogenize 3 x 4 transformation matrices."""
-        return np.concatenate(
-            [p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1
-        )
-
-    # p34_to_44 = lambda p: np.concatenate(
-    #     [p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1
-    # )
-
-    p34_to_44 = add_row_to_homogenize_transform
-
-    rays_d = poses[:, :3, 2:3]
-    rays_o = poses[:, :3, 3:4]
-
-    def min_line_dist(rays_o, rays_d):
-        A_i = np.eye(3) - rays_d * np.transpose(rays_d, [0, 2, 1])
-        b_i = -A_i @ rays_o
-        pt_mindist = np.squeeze(
-            -np.linalg.inv((np.transpose(A_i, [0, 2, 1]) @ A_i).mean(0)) @ (b_i).mean(0)
-        )
-        return pt_mindist
-
-    pt_mindist = min_line_dist(rays_o, rays_d)
-
-    center = pt_mindist
-    up = (poses[:, :3, 3] - center).mean(0)
-
-    vec0 = normalize(up)
-    vec1 = normalize(np.cross([0.1, 0.2, 0.3], vec0))
-    vec2 = normalize(np.cross(vec0, vec1))
-    pos = center
-    c2w = np.stack([vec1, vec2, vec0, pos], 1)
-
-    poses_reset = np.linalg.inv(p34_to_44(c2w[None])) @ p34_to_44(poses[:, :3, :4])
-
-    rad = np.sqrt(np.mean(np.sum(np.square(poses_reset[:, :3, 3]), -1)))
-
-    sc = 1.0 / rad
-    poses_reset[:, :3, 3] *= sc
-    bds *= sc
-    rad *= sc
-
-    centroid = np.mean(poses_reset[:, :3, 3], 0)
-    zh = centroid[2]
-    radcircle = np.sqrt(rad**2 - zh**2)
-    new_poses = []
-
-    for th in np.linspace(0.0, 2.0 * np.pi, 120):
-        camorigin = np.array([radcircle * np.cos(th), radcircle * np.sin(th), zh])
-        up = np.array([0, 0, -1.0])
-
-        vec2 = normalize(camorigin)
-        vec0 = normalize(np.cross(vec2, up))
-        vec1 = normalize(np.cross(vec2, vec0))
-        pos = camorigin
-        p = np.stack([vec0, vec1, vec2, pos], 1)
-
-        new_poses.append(p)
-
-    new_poses = np.stack(new_poses, 0)
-
-    new_poses = np.concatenate(
-        [new_poses, np.broadcast_to(poses[0, :3, -1:], new_poses[:, :3, -1:].shape)], -1
-    )
-    poses_reset = np.concatenate(
-        [
-            poses_reset[:, :3, :4],
-            np.broadcast_to(poses[0, :3, -1:], poses_reset[:, :3, -1:].shape),
-        ],
-        -1,
-    )
-
-    return poses_reset, new_poses, bds
-
-
-def _local_path(path_manager, path):
-    if path_manager is None:
-        return path
-    return path_manager.get_local_path(path)
-
-
-def _ls(path_manager, path):
-    if path_manager is None:
-        return os.listdir(path)
-    return path_manager.ls(path)
-
-
-def _exists(path_manager, path):
-    if path_manager is None:
-        return os.path.exists(path)
-    return path_manager.exists(path)
-
-
-def load_llff_data(
-    basedir,
-    factor=8,
-    recenter=True,
-    bd_factor=0.75,
-    spherify=False,
-    path_zflat=False,
-    path_manager=None,
-):
-    poses, bds, imgs = _load_data(
-        basedir, factor=factor, path_manager=path_manager
-    )  # factor=8 downsamples original imgs by 8x
-    logger.info(f"Loaded {basedir}, {bds.min()}, {bds.max()}")
-
-    # Correct rotation matrix ordering and move variable dim to axis 0
-    poses = np.concatenate([poses[:, 1:2, :], -poses[:, 0:1, :], poses[:, 2:, :]], 1)
-    poses = np.moveaxis(poses, -1, 0).astype(np.float32)
-    imgs = np.moveaxis(imgs, -1, 0).astype(np.float32)
-    images = imgs
-    bds = np.moveaxis(bds, -1, 0).astype(np.float32)
-
-    # Rescale if bd_factor is provided
-    sc = 1.0 if bd_factor is None else 1.0 / (bds.min() * bd_factor)
-    poses[:, :3, 3] *= sc
-    bds *= sc
-
-    if recenter:
-        poses = recenter_poses(poses)
-
-    if spherify:
-        poses, render_poses, bds = spherify_poses(poses, bds)
-
-    images = images.astype(np.float32)
-    poses = poses.astype(np.float32)
-
-    return images, poses, bds

pytorch3d/implicitron/dataset/orm_types.py

@@ -13,7 +13,6 @@ import struct
 from typing import Optional, Tuple
 
 import numpy as np
-
 from pytorch3d.implicitron.dataset.types import (
     DepthAnnotation,
     ImageAnnotation,
@@ -22,7 +21,6 @@ from pytorch3d.implicitron.dataset.types import (
     VideoAnnotation,
     ViewpointAnnotation,
 )
-
 from sqlalchemy import LargeBinary
 from sqlalchemy.orm import (
     composite,
@@ -92,6 +90,7 @@ def TupleTypeFactory(dtype=float, shape: Tuple[int, ...] = (2,)):
         impl = LargeBinary
         _format = format_symbol * math.prod(shape)
 
+        # pyrefly: ignore [bad-override-param-name]
         def process_bind_param(self, value, _):
             if value is None:
                 return None
@@ -101,6 +100,7 @@ def TupleTypeFactory(dtype=float, shape: Tuple[int, ...] = (2,)):
 
             return struct.pack(TupleType._format, *value)
 
+        # pyrefly: ignore [bad-override-param-name]
         def process_result_value(self, value, _):
             if value is None:
                 return None

pytorch3d/implicitron/dataset/rendered_mesh_dataset_map_provider.py

@@ -121,11 +121,17 @@ class RenderedMeshDatasetMapProvider(DatasetMapProviderBase):
         self.poses = poses.cpu()
         # pyre-ignore[16]
         self.train_dataset = SingleSceneDataset(  # pyre-ignore[28]
+            # pyrefly: ignore [unexpected-keyword]
             object_name="cow",
+            # pyrefly: ignore [unexpected-keyword]
             images=list(images.permute(0, 3, 1, 2).cpu()),
+            # pyrefly: ignore [unexpected-keyword]
             fg_probabilities=list(masks[:, None].cpu()),
+            # pyrefly: ignore [unexpected-keyword]
             poses=[self.poses[i] for i in range(len(poses))],
+            # pyrefly: ignore [unexpected-keyword]
             frame_types=[DATASET_TYPE_KNOWN] * len(poses),
+            # pyrefly: ignore [unexpected-keyword]
             eval_batches=None,
         )
 

pytorch3d/implicitron/dataset/single_sequence_dataset.py

@@ -85,7 +85,7 @@ class SingleSceneDataset(DatasetBase, Configurable):
 
 class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
     """
-    Base for provider of data for one scene from LLFF or blender datasets.
+    Base for provider of data for one scene.
 
     Members:
         base_dir: directory holding the data for the scene.
@@ -131,10 +131,11 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
             None
             # pyre-ignore[16]
             if self.fg_probabilities is None
-            else self.fg_probabilities[split]
+            else self.fg_probabilities[split]  # pyrefly: ignore [bad-index]
         )
         eval_batches = [[i] for i in range(len(split))]
         if split_idx != 0 and self.n_known_frames_for_test is not None:
+            # pyrefly: ignore [missing-attribute]
             train_split = self.i_split[0]
             if set_eval_batches:
                 generator = np.random.default_rng(seed=0)
@@ -171,40 +172,3 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
         # pyre-ignore[16]
         cameras = [self.poses[i] for i in self.i_split[0]]
         return join_cameras_as_batch(cameras)
-
-
-def _interpret_blender_cameras(
-    poses: torch.Tensor, focal: float
-) -> List[PerspectiveCameras]:
-    """
-    Convert 4x4 matrices representing cameras in blender format
-    to PyTorch3D format.
-
-    Args:
-        poses: N x 3 x 4 camera matrices
-        focal: ndc space focal length
-    """
-    pose_target_cameras = []
-    for pose_target in poses:
-        pose_target = pose_target[:3, :4]
-        mtx = torch.eye(4, dtype=pose_target.dtype)
-        mtx[:3, :3] = pose_target[:3, :3].t()
-        mtx[3, :3] = pose_target[:, 3]
-        mtx = mtx.inverse()
-
-        # flip the XZ coordinates.
-        mtx[:, [0, 2]] *= -1.0
-
-        Rpt3, Tpt3 = mtx[:, :3].split([3, 1], dim=0)
-
-        focal_length_pt3 = torch.FloatTensor([[focal, focal]])
-        principal_point_pt3 = torch.FloatTensor([[0.0, 0.0]])
-
-        cameras = PerspectiveCameras(
-            focal_length=focal_length_pt3,
-            principal_point=principal_point_pt3,
-            R=Rpt3[None],
-            T=Tpt3,
-        )
-        pose_target_cameras.append(cameras)
-    return pose_target_cameras

pytorch3d/implicitron/dataset/sql_dataset.py

@@ -10,7 +10,6 @@ import hashlib
 import json
 import logging
 import os
-
 import urllib
 from dataclasses import dataclass, Field, field
 from typing import (
@@ -32,13 +31,11 @@ import pandas as pd
 import sqlalchemy as sa
 import torch
 from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
-
 from pytorch3d.implicitron.dataset.frame_data import (
     FrameData,
     FrameDataBuilder,  # noqa
     FrameDataBuilderBase,
 )
-
 from pytorch3d.implicitron.tools.config import (
     registry,
     ReplaceableBase,
@@ -224,6 +221,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
     def __len__(self) -> int:
         return len(self._index)
 
+    # pyrefly: ignore [bad-override-param-name]
     def __getitem__(self, frame_idx: Union[int, Tuple[str, int]]) -> FrameData:
         """
         Fetches FrameData by either iloc in the index or by (sequence, frame_no) pair
@@ -427,6 +425,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
 
     # override
     @property
+    # pyrefly: ignore [bad-override]
     def frame_data_type(self) -> Type[FrameData]:
         return self.frame_data_builder.frame_data_type
 
@@ -486,9 +485,10 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
             *self._get_pick_filters(),
             *self._get_exclude_filters(),
         ]
-        if self.pick_sequences_sql_clause:
+        if pick_sequences_sql_clause := self.pick_sequences_sql_clause:
             print("Applying the custom SQL clause.")
-            where_conditions.append(sa.text(self.pick_sequences_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            where_conditions.append(sa.text(pick_sequences_sql_clause))
 
         def add_where(stmt):
             return stmt.where(*where_conditions) if where_conditions else stmt
@@ -508,6 +508,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
 
             subquery = add_where(subquery).subquery()
             stmt = sa.select(subquery.c.sequence_name).where(
+                # pyre-ignore[6]: SQLAlchemy column comparison returns ColumnElement, not bool
                 subquery.c.row_number <= self.limit_sequences_per_category_to
             )
 
@@ -631,14 +632,17 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
                 # dev load: 94 s / 23 s (3.1M / 500K)
                 pick_frames_criteria.append(
                     sa.or_(
+                        # pyrefly: ignore [missing-attribute]
                         self.frame_annotations_type._mask_mass.is_(None),
+                        # pyrefly: ignore [missing-attribute]
                         self.frame_annotations_type._mask_mass != 0,
                     )
                 )
 
-        if self.pick_frames_sql_clause:
+        if pick_frames_sql_clause := self.pick_frames_sql_clause:
             logger.info("Applying the custom SQL clause.")
-            pick_frames_criteria.append(sa.text(self.pick_frames_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            pick_frames_criteria.append(sa.text(pick_frames_sql_clause))
 
         if pick_frames_criteria:
             index = self._pick_frames_by_criteria(index, pick_frames_criteria)
@@ -695,15 +699,18 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
         if self.remove_empty_masks:
             logger.info("  excluding samples with empty masks")
             where_conditions.append(
+                # pyrefly: ignore [bad-argument-type]
                 sa.or_(
                     self.frame_annotations_type._mask_mass.is_(None),  # pyre-ignore[16]
+                    # pyrefly: ignore [missing-attribute]
                     self.frame_annotations_type._mask_mass != 0,
                 )
             )
 
-        if self.pick_frames_sql_clause:
+        if pick_frames_sql_clause := self.pick_frames_sql_clause:
             logger.info("  applying custom SQL clause")
-            where_conditions.append(sa.text(self.pick_frames_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            where_conditions.append(sa.text(pick_frames_sql_clause))
 
         if where_conditions:
             stmt = stmt.where(*where_conditions)
@@ -755,7 +762,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
         if pick_sequences:
             old_len = len(eval_batches)
             eval_batches = [b for b in eval_batches if b[0][0] in pick_sequences]
-            logger.warn(
+            logger.warning(
                 f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
             )
 
@@ -763,7 +770,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
             old_len = len(eval_batches)
             exclude_sequences = set(self.exclude_sequences)
             eval_batches = [b for b in eval_batches if b[0][0] not in exclude_sequences]
-            logger.warn(
+            logger.warning(
                 f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
             )
 

pytorch3d/implicitron/dataset/sql_dataset_provider.py

@@ -12,9 +12,7 @@ import os
 from typing import List, Optional, Tuple, Type
 
 import numpy as np
-
 from omegaconf import DictConfig, OmegaConf
-
 from pytorch3d.implicitron.dataset.dataset_map_provider import (
     DatasetMap,
     DatasetMapProviderBase,

pytorch3d/implicitron/dataset/train_eval_data_loader_provider.py

@@ -18,7 +18,6 @@ from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
 from pytorch3d.implicitron.dataset.dataset_map_provider import DatasetMap
 from pytorch3d.implicitron.dataset.frame_data import FrameData
 from pytorch3d.implicitron.tools.config import registry, run_auto_creation
-
 from torch.utils.data import DataLoader
 
 logger = logging.getLogger(__name__)

pytorch3d/implicitron/dataset/types.py

@@ -162,6 +162,7 @@ def load_dataclass(f: IO, cls: Type[_X], binary: bool = False) -> _X:
     else:
         res = _dataclass_from_dict(asdict, cls)
 
+    # pyrefly: ignore [bad-return]
     return res
 
 

pytorch3d/implicitron/dataset/utils.py

@@ -15,7 +15,6 @@ from typing import List, Optional, Tuple, TypeVar, Union
 import numpy as np
 import torch
 from PIL import Image
-
 from pytorch3d.io import IO
 from pytorch3d.renderer.cameras import PerspectiveCameras
 from pytorch3d.structures.pointclouds import Pointclouds
@@ -118,6 +117,7 @@ def crop_around_box(
     # bbox is xyxy, where the upper bound is corrected with +1
     bbox = clamp_box_to_image_bounds_and_round(
         bbox,
+        # pyrefly: ignore [bad-argument-type]
         image_size_hw=tuple(tensor.shape[-2:]),
     )
     tensor = tensor[..., bbox[1] : bbox[3], bbox[0] : bbox[2]]

pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py

@@ -312,6 +312,7 @@ def eval_batch(
             )
 
             if visualize:
+                # pyrefly: ignore [unbound-name]
                 visualizer.show_rgb(
                     results[metric_name].item(), metric_name, loss_mask_now
                 )
@@ -330,6 +331,7 @@ def eval_batch(
             results["depth_abs" + name_postfix] = abs_.mean()
 
             if visualize:
+                # pyrefly: ignore [unbound-name]
                 visualizer.show_depth(abs_.mean().item(), name_postfix, loss_mask_now)
                 if break_after_visualising:
                     breakpoint()  # noqa: B601

pytorch3d/implicitron/evaluation/evaluator.py

@@ -14,7 +14,6 @@ import warnings
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
-
 import tqdm
 from pytorch3d.implicitron.evaluation import evaluate_new_view_synthesis as evaluate
 from pytorch3d.implicitron.models.base_model import EvaluationMode, ImplicitronModelBase

pytorch3d/implicitron/models/base_model.py

@@ -10,7 +10,6 @@ from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
 
 import torch
-
 from pytorch3d.implicitron.models.renderer.base import EvaluationMode
 from pytorch3d.implicitron.tools.config import ReplaceableBase
 from pytorch3d.renderer.cameras import CamerasBase

pytorch3d/implicitron/models/generic_model.py

@@ -16,7 +16,6 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
 from omegaconf import DictConfig
-
 from pytorch3d.implicitron.models.base_model import (
     ImplicitronModelBase,
     ImplicitronRender,
@@ -28,7 +27,6 @@ from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
 )
-
 from pytorch3d.implicitron.models.renderer.base import (
     BaseRenderer,
     EvaluationMode,
@@ -38,7 +36,6 @@ from pytorch3d.implicitron.models.renderer.base import (
     RenderSamplingMode,
 )
 from pytorch3d.implicitron.models.renderer.ray_sampler import RaySamplerBase
-
 from pytorch3d.implicitron.models.utils import (
     apply_chunked,
     chunk_generator,
@@ -53,7 +50,6 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.implicitron.tools.rasterize_mc import rasterize_sparse_ray_bundle
 from pytorch3d.renderer import utils as rend_utils
 from pytorch3d.renderer.cameras import CamerasBase
@@ -476,6 +472,7 @@ class GenericModel(ImplicitronModelBase):
                 sequence_name=safe_slice_targets(sequence_name),
                 frame_timestamp=safe_slice_targets(frame_timestamp),
             )
+        # pyrefly: ignore [unsupported-operation]
         custom_args["global_code"] = global_code
 
         # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a

pytorch3d/implicitron/models/implicit_function/base.py

@@ -10,7 +10,6 @@ from abc import ABC, abstractmethod
 from typing import Optional
 
 from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
-
 from pytorch3d.implicitron.tools.config import ReplaceableBase
 from pytorch3d.renderer.cameras import CamerasBase
 

pytorch3d/implicitron/models/implicit_function/decoding_functions.py

@@ -16,14 +16,11 @@ This file contains
 
 import logging
 from dataclasses import field
-
 from enum import Enum
 from typing import Dict, Optional, Tuple
 
 import torch
-
 from omegaconf import DictConfig
-
 from pytorch3d.implicitron.tools.config import (
     Configurable,
     registry,

pytorch3d/implicitron/models/implicit_function/idr_feature_field.py

@@ -11,7 +11,6 @@ import torch
 from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
 from pytorch3d.implicitron.tools.config import registry
 from pytorch3d.renderer.implicit import HarmonicEmbedding
-
 from torch import nn
 
 from .base import ImplicitFunctionBase

pytorch3d/implicitron/models/implicit_function/neural_radiance_field.py

@@ -21,7 +21,6 @@ from pytorch3d.renderer.implicit import HarmonicEmbedding
 from pytorch3d.renderer.implicit.utils import ray_bundle_to_ray_points
 
 from .base import ImplicitFunctionBase
-
 from .decoding_functions import (  # noqa
     _xavier_init,
     MLPWithInputSkips,