Revert "Fix CUDA kernel index data type in vision/fair/pytorch3d/pytorch3d/csrc/compositing/alpha_composite.cu +10"

This reverts commit 3987612062.

.github/workflows/build.yml

@@ -21,45 +21,3 @@ 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,17 +6,6 @@ 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}/BUCK" ]]
+if [[ -f "${DIR}/TARGETS" ]]
 then
   pyfmt "${DIR}"
 else

docs/conf.py

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

docs/modules/implicitron/datasets.rst

@@ -3,6 +3,11 @@ 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:
@@ -13,6 +18,11 @@ 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

@@ -0,0 +1,56 @@
+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

@@ -0,0 +1,55 @@
+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,18 +48,22 @@ 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,
 )
@@ -193,7 +197,6 @@ 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,
         )
 
@@ -213,7 +216,6 @@ 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,6 +11,7 @@ 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,7 +14,9 @@ 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 (
@@ -169,11 +171,9 @@ 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,13 +190,8 @@ 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(
-                # pyrefly: ignore [unexpected-keyword]
-                optimizer,
-                _get_lr,
-                # pyrefly: ignore [unexpected-keyword]
-                verbose=False,
+                optimizer, _get_lr, verbose=False
             )
         else:
             raise ValueError("no such lr policy %s" % self.lr_policy)

projects/implicitron_trainer/impl/training_loop.py

@@ -199,7 +199,6 @@ 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,
@@ -216,7 +215,6 @@ 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,6 +13,13 @@ 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
@@ -84,6 +91,14 @@ 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,6 +12,7 @@ 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,7 +135,6 @@ 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.9"
+__version__ = "0.7.8"

pytorch3d/common/workaround/symeig3x3.py

@@ -82,15 +82,12 @@ class _SymEig3x3(nn.Module):
         q = inputs_trace / 3.0
 
         # Calculate squared sum of elements outside the main diagonal / 2
-        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
-        )
+        # 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)
 
         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
@@ -107,9 +104,7 @@ 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(-torch.square(p1 / (6 * self._eps))).detach()[
-            ..., None
-        ]
+        diag_soft_cond = torch.exp(-((p1 / (6 * self._eps)) ** 2)).detach()[..., None]
 
         # Eigenvalues are the ordered elements of main diagonal in the degenerate case
         diag_eigenvals, _ = torch.sort(inputs_diag, dim=-1)
@@ -175,10 +170,8 @@ 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)
@@ -206,7 +199,8 @@ class _SymEig3x3(nn.Module):
             cross_products[..., :1, :]
         )
 
-        norms_sq = torch.square(cross_products).sum(dim=-1)
+        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
+        norms_sq = (cross_products**2).sum(dim=-1)
         max_norms_index = norms_sq.argmax(dim=-1)
 
         # Pick only the cross-product with highest squared norm for each input
@@ -253,7 +247,6 @@ 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,9 +32,7 @@ __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 radius,
-    const float radius2,
-    const bool skip_points_outside_cube) {
+    const float radius2) {
   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;
@@ -53,19 +51,7 @@ __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) {
-      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
+      // Calculate the distance between the points
       scalar_t dist2 = 0.0;
       for (int d = 0; d < D; ++d) {
         scalar_t diff = p1[n][i][d] - p2[n][j][d];
@@ -91,8 +77,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
     const at::Tensor& lengths1, // (N,)
     const at::Tensor& lengths2, // (N,)
     int K,
-    float radius,
-    bool skip_points_outside_cube) {
+    float radius) {
   // 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};
@@ -135,9 +120,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
             idxs.packed_accessor64<int64_t, 3, at::RestrictPtrTraits>(),
             dists.packed_accessor64<float, 3, at::RestrictPtrTraits>(),
             K_64,
-            radius,
-            radius2,
-            skip_points_outside_cube);
+            radius2);
       }));
 
   AT_CUDA_CHECK(cudaGetLastError());

pytorch3d/csrc/ball_query/ball_query.h

@@ -25,9 +25,6 @@
 //      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
@@ -49,8 +46,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     const int K,
-    const float radius,
-    const bool skip_points_outside_cube);
+    const float radius);
 
 // CUDA implementation
 std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
@@ -59,8 +55,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCuda(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     const int K,
-    const float radius,
-    const bool skip_points_outside_cube);
+    const float radius);
 
 // Implementation which is exposed
 // Note: the backward pass reuses the KNearestNeighborBackward kernel
@@ -70,8 +65,7 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     int K,
-    float radius,
-    bool skip_points_outside_cube) {
+    float radius) {
   if (p1.is_cuda() || p2.is_cuda()) {
 #ifdef WITH_CUDA
     CHECK_CUDA(p1);
@@ -82,20 +76,16 @@ inline std::tuple<at::Tensor, at::Tensor> BallQuery(
         lengths1.contiguous(),
         lengths2.contiguous(),
         K,
-        radius,
-        skip_points_outside_cube);
+        radius);
 #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,
-      skip_points_outside_cube);
+      radius);
 }

pytorch3d/csrc/ball_query/ball_query_cpu.cpp

@@ -6,7 +6,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <math.h>
 #include <torch/extension.h>
 #include <tuple>
 
@@ -16,8 +15,7 @@ std::tuple<at::Tensor, at::Tensor> BallQueryCpu(
     const at::Tensor& lengths1,
     const at::Tensor& lengths2,
     int K,
-    float radius,
-    bool skip_points_outside_cube) {
+    float radius) {
   const int N = p1.size(0);
   const int P1 = p1.size(1);
   const int D = p1.size(2);
@@ -39,16 +37,6 @@ 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,11 +98,6 @@ 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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,9 +74,6 @@ 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);
   }
 }
@@ -104,11 +101,6 @@ 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * W * H;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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,10 +73,6 @@ 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);
   }
 }
@@ -104,11 +100,6 @@ 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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 auto batch = blockIdx.x;
+  const int batch = blockIdx.x;
 
   const int num_pixels = C * H * W;
-  const auto num_threads = gridDim.y * blockDim.x;
-  const auto tid = blockIdx.y * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.y * blockDim.x;
+  const int 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,9 +72,6 @@ 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);
   }
 }
@@ -101,11 +98,6 @@ 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,6 +8,7 @@
 
 // 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"
@@ -98,22 +99,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,8 +60,6 @@ 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);
 }
 
@@ -82,9 +80,5 @@ 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 auto tid = threadIdx.x;
+  const int 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 (auto e = blockIdx.x; e < E; e += gridDim.x) {
+  for (int 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 (auto d = tid; d < D; d += blockDim.x) {
+    for (int 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,7 +53,5 @@ 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 auto tid = threadIdx.x + blockIdx.x * blockDim.x;
-  const auto num_threads = blockDim.x * gridDim.x;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const int 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 auto tid = threadIdx.x + blockIdx.x * blockDim.x;
-  const auto num_threads = blockDim.x * gridDim.x;
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  const int 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,8 +57,6 @@ 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);
 }
 
@@ -108,9 +106,6 @@ 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,7 +44,5 @@ 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,6 +7,7 @@
  */
 
 #include <torch/extension.h>
+#include <list>
 #include <tuple>
 #include "iou_box3d/iou_utils.h"
 

pytorch3d/csrc/knn/knn.h

@@ -74,8 +74,6 @@ 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);
 }
 
@@ -142,8 +140,6 @@ 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,6 +58,5 @@ 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,8 +88,6 @@ 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);
 }
 
@@ -107,7 +105,5 @@ 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 (auto s = blockDim.x / 2; s > 32; s >>= 1) {
+    for (int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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,10 +88,6 @@ 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);
 }
@@ -147,10 +143,6 @@ 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);
 }
@@ -229,10 +221,6 @@ 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);
 }
@@ -289,10 +277,6 @@ 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);
 }
@@ -362,10 +346,6 @@ 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);
 }
@@ -416,10 +396,6 @@ 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);
 }
 
@@ -488,10 +464,6 @@ 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);
 }
@@ -542,10 +514,6 @@ 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);
 }
 
@@ -599,8 +567,6 @@ torch::Tensor PointFaceArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  CHECK_CPU(points);
-  CHECK_CPU(tris);
   return PointFaceArrayDistanceForwardCpu(points, tris, min_triangle_area);
 }
 
@@ -647,9 +613,6 @@ 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);
 }
@@ -698,8 +661,6 @@ torch::Tensor PointEdgeArrayDistanceForward(
     AT_ERROR("Not compiled with GPU support.");
 #endif
   }
-  CHECK_CPU(points);
-  CHECK_CPU(segms);
   return PointEdgeArrayDistanceForwardCpu(points, segms);
 }
 
@@ -742,8 +703,5 @@ 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,12 +104,6 @@ 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,
@@ -189,14 +183,6 @@ 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)
-using uint = unsigned int;
-using ushort = unsigned short;
+#define uint unsigned int
+#define ushort unsigned short
 #endif
 
 #include "./logging.h" // <- include before torch/extension.h

pytorch3d/csrc/pulsar/gpu/commands.h

@@ -125,45 +125,29 @@ 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;
@@ -217,16 +201,8 @@ __device__ static float atomicMin(float* address, float val) {
   ATOMICADD(&((PTR)->x), VAL.x); \
   ATOMICADD(&((PTR)->y), VAL.y); \
   ATOMICADD(&((PTR)->z), VAL.z);
-#if !defined(USE_ROCM) && (CUDART_VERSION >= 10000) && (__CUDA_ARCH__ >= 600)
+#if (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
@@ -441,7 +417,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,6 +70,11 @@ 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,6 +149,11 @@ 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(
-        self->ids_sorted_d,
-        self->n_grad_contributions_d,
+        (IntWrapper*)(self->ids_sorted_d),
+        (IntWrapper*)(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,89 +18,68 @@ 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,9 +64,8 @@ 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,9 +139,8 @@ 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;
     }
@@ -191,22 +190,20 @@ 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.
@@ -293,9 +290,8 @@ 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();
@@ -303,9 +299,8 @@ 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",
@@ -391,9 +386,8 @@ 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,10 +7,13 @@
  */
 
 #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
 
@@ -857,9 +860,8 @@ 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,9 +128,8 @@ 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,6 +6,9 @@
  * 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 (auto i = threadIdx.x; i < H * W * D; i += blockDim.x) {
+    for (int 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 auto tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const auto num_threads = blockDim.x * gridDim.x;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int 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 auto tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const auto num_threads = blockDim.x * gridDim.x;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int 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 (auto chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x) {
+  for (int 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 (auto e = threadIdx.x; e < chunk_size; e += blockDim.x) {
+    for (int 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 (auto byx = threadIdx.x; byx < num_bins_y * num_bins_x;
+    for (int 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
-  auto num_threads = gridDim.x * blockDim.x;
-  auto tid = blockDim.x * blockIdx.x + threadIdx.x;
+  int num_threads = gridDim.x * blockDim.x;
+  int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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;
-  auto num_threads = gridDim.x * blockDim.x;
-  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int num_threads = gridDim.x * blockDim.x;
+  int 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,9 +138,6 @@ 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,
@@ -235,11 +232,6 @@ 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,
@@ -314,9 +306,6 @@ 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,
@@ -434,8 +423,6 @@ 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,6 +8,7 @@
 
 #include <torch/extension.h>
 #include <algorithm>
+#include <list>
 #include <thread>
 #include <tuple>
 #include "ATen/core/TensorAccessor.h"
@@ -105,8 +106,6 @@ 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 {
@@ -119,6 +118,7 @@ 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockDim.x * blockIdx.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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 auto num_threads = gridDim.x * blockDim.x;
-  const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int num_threads = gridDim.x * blockDim.x;
+  const int 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.
-  auto num_threads = gridDim.x * blockDim.x;
-  auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int num_threads = gridDim.x * blockDim.x;
+  int 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,10 +91,6 @@ 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,
@@ -170,10 +166,6 @@ 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,
@@ -240,8 +232,6 @@ 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");
   }
 }
@@ -294,10 +284,6 @@ 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,8 +107,7 @@ 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 int64_t max_K_known = -1) {
+    const at::Tensor& start_idxs) {
   // 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};
@@ -130,12 +129,7 @@ at::Tensor FarthestPointSamplingCuda(
 
   const int64_t N = points.size(0);
   const int64_t P = points.size(1);
-  int64_t max_K;
-  if (max_K_known > 0) {
-    max_K = max_K_known;
-  } else {
-    max_K = at::max(K).item<int64_t>();
-  }
+  const int64_t 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,8 +43,7 @@ at::Tensor FarthestPointSamplingCuda(
     const at::Tensor& points,
     const at::Tensor& lengths,
     const at::Tensor& K,
-    const at::Tensor& start_idxs,
-    const int64_t max_K_known = -1);
+    const at::Tensor& start_idxs);
 
 at::Tensor FarthestPointSamplingCpu(
     const at::Tensor& points,
@@ -57,23 +56,17 @@ at::Tensor FarthestPointSampling(
     const at::Tensor& points,
     const at::Tensor& lengths,
     const at::Tensor& K,
-    const at::Tensor& start_idxs,
-    const int64_t max_K_known = -1) {
+    const at::Tensor& start_idxs) {
   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, max_K_known);
+    return FarthestPointSamplingCuda(points, lengths, K, start_idxs);
 #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,8 +71,6 @@ 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,7 +99,8 @@ 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,
@@ -123,7 +124,8 @@ 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,
@@ -143,7 +145,8 @@ 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,
@@ -200,7 +203,8 @@ 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,
@@ -239,7 +243,8 @@ 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,
@@ -278,7 +283,8 @@ 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,
@@ -307,7 +313,8 @@ 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,
@@ -323,7 +330,8 @@ 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,7 +15,3 @@
 #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;
-  using scalar_t = T;
+  typedef T scalar_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;
-  using scalar_t = T;
+  typedef T scalar_t;
   vec3(T x, T y, T z) : x(x), y(y), z(z) {}
 };
 

pytorch3d/csrc/utils/warp_reduce.cuh

@@ -10,20 +10,10 @@
 #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,
-// 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.
+// 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.
 
 template <typename scalar_t>
 __device__ void
@@ -33,6 +23,8 @@ 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,7 +200,6 @@ 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.
@@ -371,7 +370,6 @@ 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,10 +62,8 @@ 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(
@@ -77,14 +75,10 @@ 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(
@@ -95,10 +89,8 @@ 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(
@@ -466,8 +458,6 @@ 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,7 +52,6 @@ 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
@@ -148,17 +147,12 @@ 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,14 +34,12 @@ 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

@@ -0,0 +1,55 @@
+# 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,12 +64,16 @@ 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,6 +21,7 @@ from typing import (
 )
 
 import torch
+
 from pytorch3d.implicitron.dataset.frame_data import FrameData
 from pytorch3d.implicitron.dataset.utils import GenericWorkaround
 
@@ -101,7 +102,6 @@ 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,6 +25,7 @@ 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_,
@@ -297,11 +298,7 @@ class FrameData(Mapping[str, Any]):
         depth_map = self.depth_map
         if depth_map is not None:
             clamp_bbox_xyxy_depth = rescale_bbox(
-                # pyrefly: ignore [bad-argument-type]
-                clamp_bbox_xyxy,
-                # pyrefly: ignore [bad-argument-type]
-                tuple(depth_map.shape[-2:]),
-                effective_image_size_hw,
+                clamp_bbox_xyxy, tuple(depth_map.shape[-2:]), effective_image_size_hw
             ).long()
             self.depth_map = crop_around_box(
                 depth_map,
@@ -312,11 +309,7 @@ class FrameData(Mapping[str, Any]):
         depth_mask = self.depth_mask
         if depth_mask is not None:
             clamp_bbox_xyxy_depth = rescale_bbox(
-                # pyrefly: ignore [bad-argument-type]
-                clamp_bbox_xyxy,
-                # pyrefly: ignore [bad-argument-type]
-                tuple(depth_mask.shape[-2:]),
-                effective_image_size_hw,
+                clamp_bbox_xyxy, tuple(depth_mask.shape[-2:]), effective_image_size_hw
             ).long()
             self.depth_mask = crop_around_box(
                 depth_mask,
@@ -461,7 +454,6 @@ 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,6 +38,7 @@ 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
 
 
@@ -326,9 +327,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 = [
@@ -425,7 +426,6 @@ 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,6 +21,7 @@ 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,6 +18,7 @@ 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,
@@ -30,6 +31,7 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
+
 from pytorch3d.renderer.cameras import CamerasBase
 from tqdm import tqdm
 
@@ -311,7 +313,6 @@ class JsonIndexDatasetMapProviderV2(DatasetMapProviderBase):
                     subset_mapping["test"],
                 ) = self._extend_test_data_with_known_views(
                     subset_mapping,
-                    # pyrefly: ignore [bad-argument-type]
                     eval_batch_index,
                 )
 
@@ -323,7 +324,6 @@ 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,7 +337,6 @@ 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

@@ -0,0 +1,69 @@
+# 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

@@ -0,0 +1,143 @@
+# @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

@@ -0,0 +1,336 @@
+# @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,6 +13,7 @@ import struct
 from typing import Optional, Tuple
 
 import numpy as np
+
 from pytorch3d.implicitron.dataset.types import (
     DepthAnnotation,
     ImageAnnotation,
@@ -21,6 +22,7 @@ from pytorch3d.implicitron.dataset.types import (
     VideoAnnotation,
     ViewpointAnnotation,
 )
+
 from sqlalchemy import LargeBinary
 from sqlalchemy.orm import (
     composite,
@@ -90,7 +92,6 @@ 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
@@ -100,7 +101,6 @@ 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,17 +121,11 @@ 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.
+    Base for provider of data for one scene from LLFF or blender datasets.
 
     Members:
         base_dir: directory holding the data for the scene.
@@ -131,11 +131,10 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
             None
             # pyre-ignore[16]
             if self.fg_probabilities is None
-            else self.fg_probabilities[split]  # pyrefly: ignore [bad-index]
+            else self.fg_probabilities[split]
         )
         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)
@@ -172,3 +171,40 @@ 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,6 +10,7 @@ import hashlib
 import json
 import logging
 import os
+
 import urllib
 from dataclasses import dataclass, Field, field
 from typing import (
@@ -31,11 +32,13 @@ 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,
@@ -221,7 +224,6 @@ 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
@@ -425,7 +427,6 @@ 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
 
@@ -485,10 +486,9 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
             *self._get_pick_filters(),
             *self._get_exclude_filters(),
         ]
-        if pick_sequences_sql_clause := self.pick_sequences_sql_clause:
+        if self.pick_sequences_sql_clause:
             print("Applying the custom SQL clause.")
-            # pyre-ignore[6]: TextClause is compatible with where conditions
-            where_conditions.append(sa.text(pick_sequences_sql_clause))
+            where_conditions.append(sa.text(self.pick_sequences_sql_clause))
 
         def add_where(stmt):
             return stmt.where(*where_conditions) if where_conditions else stmt
@@ -508,7 +508,6 @@ 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
             )
 
@@ -632,17 +631,14 @@ 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 pick_frames_sql_clause := self.pick_frames_sql_clause:
+        if self.pick_frames_sql_clause:
             logger.info("Applying the custom SQL clause.")
-            # pyre-ignore[6]: TextClause is compatible with where conditions
-            pick_frames_criteria.append(sa.text(pick_frames_sql_clause))
+            pick_frames_criteria.append(sa.text(self.pick_frames_sql_clause))
 
         if pick_frames_criteria:
             index = self._pick_frames_by_criteria(index, pick_frames_criteria)
@@ -699,18 +695,15 @@ 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 pick_frames_sql_clause := self.pick_frames_sql_clause:
+        if self.pick_frames_sql_clause:
             logger.info("  applying custom SQL clause")
-            # pyre-ignore[6]: TextClause is compatible with where conditions
-            where_conditions.append(sa.text(pick_frames_sql_clause))
+            where_conditions.append(sa.text(self.pick_frames_sql_clause))
 
         if where_conditions:
             stmt = stmt.where(*where_conditions)
@@ -762,7 +755,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.warning(
+            logger.warn(
                 f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
             )
 
@@ -770,7 +763,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.warning(
+            logger.warn(
                 f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
             )
 

pytorch3d/implicitron/dataset/sql_dataset_provider.py

@@ -12,7 +12,9 @@ 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,6 +18,7 @@ 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,7 +162,6 @@ 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,6 +15,7 @@ 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
@@ -117,7 +118,6 @@ 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,7 +312,6 @@ def eval_batch(
             )
 
             if visualize:
-                # pyrefly: ignore [unbound-name]
                 visualizer.show_rgb(
                     results[metric_name].item(), metric_name, loss_mask_now
                 )
@@ -331,7 +330,6 @@ 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,6 +14,7 @@ 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,6 +10,7 @@ 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,6 +16,7 @@ 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,
@@ -27,6 +28,7 @@ from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
 )
+
 from pytorch3d.implicitron.models.renderer.base import (
     BaseRenderer,
     EvaluationMode,
@@ -36,6 +38,7 @@ 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,
@@ -50,6 +53,7 @@ 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
@@ -472,7 +476,6 @@ 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,6 +10,7 @@ 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,11 +16,14 @@ 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,6 +11,7 @@ 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,6 +21,7 @@ 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,