fbcode/vision/fair/pytorch3d/pytorch3d/transforms/rotation_conversions.py

Reviewed By: bottler

Differential Revision: D93712828

fbshipit-source-id: 3465af450104bb1e5f491e3c0ee0259698cf8ceb

dev/linter.sh

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

docs/conf.py

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

projects/implicitron_trainer/experiment.py

@@ -48,22 +48,18 @@ The outputs of the experiment are saved and logged in multiple ways:
 import logging
 import os
 import warnings
-
 from dataclasses import field
 
 import hydra
-
 import torch
 from accelerate import Accelerator
 from omegaconf import DictConfig, OmegaConf
 from packaging import version
-
 from pytorch3d.implicitron.dataset.data_source import (
     DataSourceBase,
     ImplicitronDataSource,
 )
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
-
 from pytorch3d.implicitron.models.renderer.multipass_ea import (
     MultiPassEmissionAbsorptionRenderer,
 )

projects/implicitron_trainer/impl/model_factory.py

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

projects/implicitron_trainer/impl/optimizer_factory.py

@@ -14,9 +14,7 @@ from dataclasses import field
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch.optim
-
 from accelerate import Accelerator
-
 from pytorch3d.implicitron.models.base_model import ImplicitronModelBase
 from pytorch3d.implicitron.tools import model_io
 from pytorch3d.implicitron.tools.config import (

projects/implicitron_trainer/tests/test_experiment.py

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

pytorch3d/__init__.py

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

pytorch3d/csrc/ball_query/ball_query.cu

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

pytorch3d/csrc/ball_query/ball_query.h

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

pytorch3d/csrc/ball_query/ball_query_cpu.cpp

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

pytorch3d/csrc/blending/sigmoid_alpha_blend.h

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

pytorch3d/csrc/compositing/alpha_composite.h

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

pytorch3d/csrc/compositing/norm_weighted_sum.h

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

pytorch3d/csrc/compositing/weighted_sum.h

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

pytorch3d/csrc/ext.cpp

@@ -8,7 +8,6 @@
 
 // clang-format off
 #include "./pulsar/global.h" // Include before <torch/extension.h>.
-#include <torch/extension.h>
 // clang-format on
 #include "./pulsar/pytorch/renderer.h"
 #include "./pulsar/pytorch/tensor_util.h"
@@ -106,15 +105,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   py::class_<
       pulsar::pytorch::Renderer,
       std::shared_ptr<pulsar::pytorch::Renderer>>(m, "PulsarRenderer")
-      .def(py::init<
-           const uint&,
-           const uint&,
-           const uint&,
-           const bool&,
-           const bool&,
-           const float&,
-           const uint&,
-           const uint&>())
+      .def(
+          py::init<
+              const uint&,
+              const uint&,
+              const uint&,
+              const bool&,
+              const bool&,
+              const float&,
+              const uint&,
+              const uint&>())
       .def(
           "__eq__",
           [](const pulsar::pytorch::Renderer& a,

pytorch3d/csrc/face_areas_normals/face_areas_normals.h

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

pytorch3d/csrc/gather_scatter/gather_scatter.h

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

pytorch3d/csrc/interp_face_attrs/interp_face_attrs.h

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

pytorch3d/csrc/iou_box3d/iou_box3d.h

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

pytorch3d/csrc/knn/knn.h

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

pytorch3d/csrc/marching_cubes/marching_cubes.h

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

pytorch3d/csrc/packed_to_padded_tensor/packed_to_padded_tensor.h

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

pytorch3d/csrc/point_mesh/point_mesh_cpu.cpp

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

pytorch3d/csrc/point_mesh/point_mesh_cuda.h

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

pytorch3d/csrc/points_to_volumes/points_to_volumes.h

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

pytorch3d/csrc/pulsar/global.h

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

pytorch3d/csrc/pulsar/gpu/commands.h

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

pytorch3d/csrc/pulsar/include/math.h

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

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

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

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

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

pytorch3d/csrc/pulsar/include/renderer.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/warnings.cpp

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

pytorch3d/csrc/rasterize_meshes/rasterize_meshes.h

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

pytorch3d/csrc/rasterize_points/rasterize_points.h

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

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.cu

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

pytorch3d/csrc/sample_farthest_points/sample_farthest_points.h

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

pytorch3d/csrc/sample_pdf/sample_pdf.h

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

pytorch3d/csrc/utils/dispatch.cuh

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

pytorch3d/csrc/utils/pytorch3d_cutils.h

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

pytorch3d/implicitron/dataset/dataset_base.py

@@ -21,7 +21,6 @@ from typing import (
 )
 
 import torch
-
 from pytorch3d.implicitron.dataset.frame_data import FrameData
 from pytorch3d.implicitron.dataset.utils import GenericWorkaround
 

pytorch3d/implicitron/dataset/frame_data.py

@@ -25,7 +25,6 @@ from typing import (
 
 import numpy as np
 import torch
-
 from pytorch3d.implicitron.dataset import orm_types, types
 from pytorch3d.implicitron.dataset.utils import (
     adjust_camera_to_bbox_crop_,

pytorch3d/implicitron/dataset/json_index_dataset.py

@@ -38,7 +38,6 @@ from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar
 from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
 from pytorch3d.renderer.camera_utils import join_cameras_as_batch
 from pytorch3d.renderer.cameras import CamerasBase
-
 from tqdm import tqdm
 
 
@@ -327,9 +326,9 @@ class JsonIndexDataset(DatasetBase, ReplaceableBase):
                 assert os.path.normpath(
                     # pyre-ignore[16]
                     self.frame_annots[idx]["frame_annotation"].image.path
-                ) == os.path.normpath(
-                    path
-                ), f"Inconsistent frame indices {seq_name, frame_no, path}."
+                ) == os.path.normpath(path), (
+                    f"Inconsistent frame indices {seq_name, frame_no, path}."
+                )
             return idx
 
         dataset_idx = [

pytorch3d/implicitron/dataset/json_index_dataset_map_provider.py

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

pytorch3d/implicitron/dataset/json_index_dataset_map_provider_v2.py

@@ -18,7 +18,6 @@ from typing import Dict, List, Optional, Tuple, Type, Union
 
 import numpy as np
 from iopath.common.file_io import PathManager
-
 from omegaconf import DictConfig
 from pytorch3d.implicitron.dataset.dataset_map_provider import (
     DatasetMap,
@@ -31,7 +30,6 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.renderer.cameras import CamerasBase
 from tqdm import tqdm
 

pytorch3d/implicitron/dataset/llff_dataset_map_provider.py

@@ -12,7 +12,6 @@ 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,

pytorch3d/implicitron/dataset/load_llff.py

@@ -8,7 +8,6 @@ import os
 import warnings
 
 import numpy as np
-
 from PIL import Image
 
 

pytorch3d/implicitron/dataset/orm_types.py

@@ -13,7 +13,6 @@ import struct
 from typing import Optional, Tuple
 
 import numpy as np
-
 from pytorch3d.implicitron.dataset.types import (
     DepthAnnotation,
     ImageAnnotation,
@@ -22,7 +21,6 @@ from pytorch3d.implicitron.dataset.types import (
     VideoAnnotation,
     ViewpointAnnotation,
 )
-
 from sqlalchemy import LargeBinary
 from sqlalchemy.orm import (
     composite,

pytorch3d/implicitron/dataset/sql_dataset.py

@@ -10,7 +10,6 @@ import hashlib
 import json
 import logging
 import os
-
 import urllib
 from dataclasses import dataclass, Field, field
 from typing import (
@@ -32,13 +31,11 @@ import pandas as pd
 import sqlalchemy as sa
 import torch
 from pytorch3d.implicitron.dataset.dataset_base import DatasetBase
-
 from pytorch3d.implicitron.dataset.frame_data import (
     FrameData,
     FrameDataBuilder,  # noqa
     FrameDataBuilderBase,
 )
-
 from pytorch3d.implicitron.tools.config import (
     registry,
     ReplaceableBase,
@@ -486,9 +483,10 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
             *self._get_pick_filters(),
             *self._get_exclude_filters(),
         ]
-        if self.pick_sequences_sql_clause:
+        if pick_sequences_sql_clause := self.pick_sequences_sql_clause:
             print("Applying the custom SQL clause.")
-            where_conditions.append(sa.text(self.pick_sequences_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            where_conditions.append(sa.text(pick_sequences_sql_clause))
 
         def add_where(stmt):
             return stmt.where(*where_conditions) if where_conditions else stmt
@@ -508,6 +506,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
 
             subquery = add_where(subquery).subquery()
             stmt = sa.select(subquery.c.sequence_name).where(
+                # pyre-ignore[6]: SQLAlchemy column comparison returns ColumnElement, not bool
                 subquery.c.row_number <= self.limit_sequences_per_category_to
             )
 
@@ -636,9 +635,10 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
                     )
                 )
 
-        if self.pick_frames_sql_clause:
+        if pick_frames_sql_clause := self.pick_frames_sql_clause:
             logger.info("Applying the custom SQL clause.")
-            pick_frames_criteria.append(sa.text(self.pick_frames_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            pick_frames_criteria.append(sa.text(pick_frames_sql_clause))
 
         if pick_frames_criteria:
             index = self._pick_frames_by_criteria(index, pick_frames_criteria)
@@ -701,9 +701,10 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
                 )
             )
 
-        if self.pick_frames_sql_clause:
+        if pick_frames_sql_clause := self.pick_frames_sql_clause:
             logger.info("  applying custom SQL clause")
-            where_conditions.append(sa.text(self.pick_frames_sql_clause))
+            # pyre-ignore[6]: TextClause is compatible with where conditions
+            where_conditions.append(sa.text(pick_frames_sql_clause))
 
         if where_conditions:
             stmt = stmt.where(*where_conditions)
@@ -755,7 +756,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
         if pick_sequences:
             old_len = len(eval_batches)
             eval_batches = [b for b in eval_batches if b[0][0] in pick_sequences]
-            logger.warn(
+            logger.warning(
                 f"Picked eval batches by sequence/cat: {old_len} -> {len(eval_batches)}"
             )
 
@@ -763,7 +764,7 @@ class SqlIndexDataset(DatasetBase, ReplaceableBase):
             old_len = len(eval_batches)
             exclude_sequences = set(self.exclude_sequences)
             eval_batches = [b for b in eval_batches if b[0][0] not in exclude_sequences]
-            logger.warn(
+            logger.warning(
                 f"Excluded eval batches by sequence: {old_len} -> {len(eval_batches)}"
             )
 

pytorch3d/implicitron/dataset/sql_dataset_provider.py

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

pytorch3d/implicitron/dataset/train_eval_data_loader_provider.py

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

pytorch3d/implicitron/dataset/utils.py

@@ -15,7 +15,6 @@ from typing import List, Optional, Tuple, TypeVar, Union
 import numpy as np
 import torch
 from PIL import Image
-
 from pytorch3d.io import IO
 from pytorch3d.renderer.cameras import PerspectiveCameras
 from pytorch3d.structures.pointclouds import Pointclouds

pytorch3d/implicitron/evaluation/evaluator.py

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

pytorch3d/implicitron/models/base_model.py

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

pytorch3d/implicitron/models/generic_model.py

@@ -16,7 +16,6 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
 from omegaconf import DictConfig
-
 from pytorch3d.implicitron.models.base_model import (
     ImplicitronModelBase,
     ImplicitronRender,
@@ -28,7 +27,6 @@ from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
 )
-
 from pytorch3d.implicitron.models.renderer.base import (
     BaseRenderer,
     EvaluationMode,
@@ -38,7 +36,6 @@ from pytorch3d.implicitron.models.renderer.base import (
     RenderSamplingMode,
 )
 from pytorch3d.implicitron.models.renderer.ray_sampler import RaySamplerBase
-
 from pytorch3d.implicitron.models.utils import (
     apply_chunked,
     chunk_generator,
@@ -53,7 +50,6 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.implicitron.tools.rasterize_mc import rasterize_sparse_ray_bundle
 from pytorch3d.renderer import utils as rend_utils
 from pytorch3d.renderer.cameras import CamerasBase

pytorch3d/implicitron/models/implicit_function/base.py

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

pytorch3d/implicitron/models/implicit_function/decoding_functions.py

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

pytorch3d/implicitron/models/implicit_function/idr_feature_field.py

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

pytorch3d/implicitron/models/implicit_function/neural_radiance_field.py

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

pytorch3d/implicitron/models/implicit_function/utils.py

@@ -9,7 +9,6 @@
 from typing import Callable, Optional
 
 import torch
-
 import torch.nn.functional as F
 from pytorch3d.common.compat import prod
 from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle

pytorch3d/implicitron/models/implicit_function/voxel_grid.py

@@ -21,8 +21,6 @@ import logging
 import warnings
 from collections.abc import Mapping
 from dataclasses import dataclass, field
-
-from distutils.version import LooseVersion
 from typing import Any, Callable, ClassVar, Dict, Iterator, List, Optional, Tuple, Type
 
 import torch
@@ -222,7 +220,8 @@ class VoxelGridBase(ReplaceableBase, torch.nn.Module):
                 + "| 'bicubic' | 'linear' | 'area' | 'nearest-exact'"
             )
 
-        interpolate_has_antialias = LooseVersion(torch.__version__) >= "1.11"
+        # We assume PyTorch 1.11 and newer.
+        interpolate_has_antialias = True
 
         if antialias and not interpolate_has_antialias:
             warnings.warn("Antialiased interpolation requires PyTorch 1.11+; ignoring")

pytorch3d/implicitron/models/implicit_function/voxel_grid_implicit_function.py

@@ -13,9 +13,7 @@ from dataclasses import fields
 from typing import Callable, Dict, Optional, Tuple
 
 import torch
-
 from omegaconf import DictConfig
-
 from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase
 from pytorch3d.implicitron.models.implicit_function.decoding_functions import (
     DecoderFunctionBase,

pytorch3d/implicitron/models/overfit_model.py

@@ -17,7 +17,6 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING, Un
 
 import torch
 from omegaconf import DictConfig
-
 from pytorch3d.implicitron.models.base_model import (
     ImplicitronModelBase,
     ImplicitronRender,
@@ -28,7 +27,6 @@ from pytorch3d.implicitron.models.metrics import (
     RegularizationMetricsBase,
     ViewMetricsBase,
 )
-
 from pytorch3d.implicitron.models.renderer.base import (
     BaseRenderer,
     EvaluationMode,
@@ -50,7 +48,6 @@ from pytorch3d.implicitron.tools.config import (
     registry,
     run_auto_creation,
 )
-
 from pytorch3d.implicitron.tools.rasterize_mc import rasterize_sparse_ray_bundle
 from pytorch3d.renderer import utils as rend_utils
 from pytorch3d.renderer.cameras import CamerasBase

pytorch3d/implicitron/models/renderer/ray_point_refiner.py

@@ -11,7 +11,6 @@ import copy
 import torch
 from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
 from pytorch3d.implicitron.tools.config import Configurable, expand_args_fields
-
 from pytorch3d.renderer.implicit.sample_pdf import sample_pdf
 
 

pytorch3d/implicitron/models/renderer/rgb_net.py

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

pytorch3d/implicitron/models/utils.py

@@ -17,11 +17,8 @@ from typing import Any, Dict, Optional, Tuple
 import torch
 import tqdm
 from pytorch3d.common.compat import prod
-
 from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
-
 from pytorch3d.implicitron.tools import image_utils
-
 from pytorch3d.implicitron.tools.utils import cat_dataclass
 
 
@@ -83,9 +80,9 @@ def preprocess_input(
 
     if mask_depths and fg_mask is not None and depth_map is not None:
         # mask the depths
-        assert (
-            mask_threshold > 0.0
-        ), "Depths should be masked only with thresholded masks"
+        assert mask_threshold > 0.0, (
+            "Depths should be masked only with thresholded masks"
+        )
         warnings.warn("Masking depths!")
         depth_map = depth_map * fg_mask
 

pytorch3d/implicitron/models/visualization/render_flyaround.py

@@ -304,11 +304,11 @@ def _show_predictions(
     assert isinstance(preds, list)
 
     pred_all = []
-    # Randomly choose a subset of the rendered images, sort by ordr in the sequence
+    # Randomly choose a subset of the rendered images, sort by order in the sequence
     n_samples = min(n_samples, len(preds))
     pred_idx = sorted(random.sample(list(range(len(preds))), n_samples))
     for predi in pred_idx:
-        # Make the concatentation for the same camera vertically
+        # Make the concatenation for the same camera vertically
         pred_all.append(
             torch.cat(
                 [
@@ -359,7 +359,7 @@ def _generate_prediction_videos(
     vws = {}
     for k in predicted_keys:
         if k not in preds[0]:
-            logger.warn(f"Cannot generate video for prediction key '{k}'")
+            logger.warning(f"Cannot generate video for prediction key '{k}'")
             continue
         cache_dir = (
             None

pytorch3d/implicitron/tools/rasterize_mc.py

@@ -10,7 +10,6 @@ import math
 from typing import Optional, Tuple
 
 import pytorch3d
-
 import torch
 from pytorch3d.ops import packed_to_padded
 from pytorch3d.renderer import PerspectiveCameras

pytorch3d/implicitron/tools/stats.py

@@ -499,7 +499,7 @@ class StatsJSONEncoder(json.JSONEncoder):
             return enc
         else:
             raise TypeError(
-                f"Object of type {o.__class__.__name__} " f"is not JSON serializable"
+                f"Object of type {o.__class__.__name__} is not JSON serializable"
             )
 
 

pytorch3d/implicitron/tools/video_writer.py

@@ -17,7 +17,6 @@ import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-
 from PIL import Image
 
 _NO_TORCHVISION = False

pytorch3d/io/obj_io.py

@@ -796,7 +796,7 @@ def save_obj(
         # Create .mtl file with the material name and texture map filename
         # TODO: enable material properties to also be saved.
         with _open_file(mtl_path, path_manager, "w") as f_mtl:
-            lines = f"newmtl mesh\n" f"map_Kd {output_path.stem}.png\n"
+            lines = f"newmtl mesh\nmap_Kd {output_path.stem}.png\n"
             f_mtl.write(lines)
 
 

pytorch3d/loss/__init__.py

@@ -8,11 +8,8 @@
 
 
 from .chamfer import chamfer_distance
-
 from .mesh_edge_loss import mesh_edge_loss
-
 from .mesh_laplacian_smoothing import mesh_laplacian_smoothing
-
 from .mesh_normal_consistency import mesh_normal_consistency
 from .point_mesh_distance import point_mesh_edge_distance, point_mesh_face_distance
 

pytorch3d/loss/mesh_laplacian_smoothing.py

@@ -114,9 +114,7 @@ def mesh_laplacian_smoothing(meshes, method: str = "uniform"):
             if method == "cot":
                 norm_w = torch.sparse.sum(L, dim=1).to_dense().view(-1, 1)
                 idx = norm_w > 0
-                # pyre-fixme[58]: `/` is not supported for operand types `float` and
-                #  `Tensor`.
-                norm_w[idx] = 1.0 / norm_w[idx]
+                norm_w[idx] = torch.reciprocal(norm_w[idx])
             else:
                 L_sum = torch.sparse.sum(L, dim=1).to_dense().view(-1, 1)
                 norm_w = 0.25 * inv_areas

pytorch3d/loss/point_mesh_distance.py

@@ -6,6 +6,7 @@
 
 # pyre-unsafe
 
+import torch
 from pytorch3d import _C
 from pytorch3d.structures import Meshes, Pointclouds
 from torch.autograd import Function
@@ -302,8 +303,7 @@ def point_mesh_edge_distance(meshes: Meshes, pcls: Pointclouds):
     point_to_cloud_idx = pcls.packed_to_cloud_idx()  # (sum(P_i), )
     num_points_per_cloud = pcls.num_points_per_cloud()  # (N,)
     weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    weights_p = 1.0 / weights_p.float()
+    weights_p = torch.reciprocal(weights_p.float())
     point_to_edge = point_to_edge * weights_p
     point_dist = point_to_edge.sum() / N
 
@@ -377,8 +377,7 @@ def point_mesh_face_distance(
     point_to_cloud_idx = pcls.packed_to_cloud_idx()  # (sum(P_i),)
     num_points_per_cloud = pcls.num_points_per_cloud()  # (N,)
     weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    weights_p = 1.0 / weights_p.float()
+    weights_p = torch.reciprocal(weights_p.float())
     point_to_face = point_to_face * weights_p
     point_dist = point_to_face.sum() / N
 

pytorch3d/ops/__init__.py

@@ -8,17 +8,14 @@
 
 from .ball_query import ball_query
 from .cameras_alignment import corresponding_cameras_alignment
-
 from .cubify import cubify
 from .graph_conv import GraphConv
 from .interp_face_attrs import interpolate_face_attributes
 from .iou_box3d import box3d_overlap
 from .knn import knn_gather, knn_points
 from .laplacian_matrices import cot_laplacian, laplacian, norm_laplacian
-
 from .mesh_face_areas_normals import mesh_face_areas_normals
 from .mesh_filtering import taubin_smoothing
-
 from .packed_to_padded import packed_to_padded, padded_to_packed
 from .perspective_n_points import efficient_pnp
 from .points_alignment import corresponding_points_alignment, iterative_closest_point
@@ -30,9 +27,7 @@ from .points_to_volumes import (
     add_pointclouds_to_volumes,
     add_points_features_to_volume_densities_features,
 )
-
 from .sample_farthest_points import sample_farthest_points
-
 from .sample_points_from_meshes import sample_points_from_meshes
 from .subdivide_meshes import SubdivideMeshes
 from .utils import (
@@ -42,7 +37,6 @@ from .utils import (
     is_pointclouds,
     wmean,
 )
-
 from .vert_align import vert_align
 
 

pytorch3d/ops/ball_query.py

@@ -23,11 +23,13 @@ class _ball_query(Function):
     """
 
     @staticmethod
-    def forward(ctx, p1, p2, lengths1, lengths2, K, radius):
+    def forward(ctx, p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube):
         """
         Arguments defintions the same as in the ball_query function
         """
-        idx, dists = _C.ball_query(p1, p2, lengths1, lengths2, K, radius)
+        idx, dists = _C.ball_query(
+            p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube
+        )
         ctx.save_for_backward(p1, p2, lengths1, lengths2, idx)
         ctx.mark_non_differentiable(idx)
         return dists, idx
@@ -49,7 +51,7 @@ class _ball_query(Function):
         grad_p1, grad_p2 = _C.knn_points_backward(
             p1, p2, lengths1, lengths2, idx, 2, grad_dists
         )
-        return grad_p1, grad_p2, None, None, None, None
+        return grad_p1, grad_p2, None, None, None, None, None
 
 
 def ball_query(
@@ -60,6 +62,7 @@ def ball_query(
     K: int = 500,
     radius: float = 0.2,
     return_nn: bool = True,
+    skip_points_outside_cube: bool = False,
 ):
     """
     Ball Query is an alternative to KNN. It can be
@@ -98,6 +101,9 @@ def ball_query(
             within the radius
         radius: the radius around each point within which the neighbors need to be located
         return_nn: If set to True returns the K neighbor points in p2 for each point in p1.
+        skip_points_outside_cube: If set to 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:
         dists: Tensor of shape (N, P1, K) giving the squared distances to
@@ -134,7 +140,9 @@ def ball_query(
     if lengths2 is None:
         lengths2 = torch.full((N,), P2, dtype=torch.int64, device=p1.device)
 
-    dists, idx = _ball_query.apply(p1, p2, lengths1, lengths2, K, radius)
+    dists, idx = _ball_query.apply(
+        p1, p2, lengths1, lengths2, K, radius, skip_points_outside_cube
+    )
 
     # Gather the neighbors if needed
     points_nn = masked_gather(p2, idx) if return_nn else None

pytorch3d/ops/cubify.py

@@ -11,9 +11,7 @@ from typing import Optional
 
 import torch
 import torch.nn.functional as F
-
 from pytorch3d.common.compat import meshgrid_ij
-
 from pytorch3d.structures import Meshes
 
 

pytorch3d/ops/laplacian_matrices.py

@@ -47,8 +47,7 @@ def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
     # i.e. A[i, j] = 1 if (i,j) is an edge, or
     # A[e0, e1] = 1 &  A[e1, e0] = 1
     ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
-    # pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
-    A = torch.sparse.FloatTensor(idx, ones, (V, V))
+    A = torch.sparse_coo_tensor(idx, ones, (V, V), dtype=torch.float32)
 
     # the sum of i-th row of A gives the degree of the i-th vertex
     deg = torch.sparse.sum(A, dim=1).to_dense()
@@ -56,21 +55,17 @@ def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
     # We construct the Laplacian matrix by adding the non diagonal values
     # i.e. L[i, j] = 1 ./ deg(i) if (i, j) is an edge
     deg0 = deg[e0]
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    deg0 = torch.where(deg0 > 0.0, 1.0 / deg0, deg0)
+    deg0 = torch.where(deg0 > 0.0, torch.reciprocal(deg0), deg0)
     deg1 = deg[e1]
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    deg1 = torch.where(deg1 > 0.0, 1.0 / deg1, deg1)
+    deg1 = torch.where(deg1 > 0.0, torch.reciprocal(deg1), deg1)
     val = torch.cat([deg0, deg1])
-    # pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
-    L = torch.sparse.FloatTensor(idx, val, (V, V))
+    L = torch.sparse_coo_tensor(idx, val, (V, V), dtype=torch.float32)
 
     # Then we add the diagonal values L[i, i] = -1.
     idx = torch.arange(V, device=verts.device)
     idx = torch.stack([idx, idx], dim=0)
     ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
-    # pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
-    L -= torch.sparse.FloatTensor(idx, ones, (V, V))
+    L -= torch.sparse_coo_tensor(idx, ones, (V, V), dtype=torch.float32)
 
     return L
 
@@ -126,8 +121,7 @@ def cot_laplacian(
     ii = faces[:, [1, 2, 0]]
     jj = faces[:, [2, 0, 1]]
     idx = torch.stack([ii, jj], dim=0).view(2, F * 3)
-    # pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
-    L = torch.sparse.FloatTensor(idx, cot.view(-1), (V, V))
+    L = torch.sparse_coo_tensor(idx, cot.view(-1), (V, V), dtype=torch.float32)
 
     # Make it symmetric; this means we are also setting
     # L[v2, v1] = cota
@@ -141,8 +135,7 @@ def cot_laplacian(
     val = torch.stack([area] * 3, dim=1).view(-1)
     inv_areas.scatter_add_(0, idx, val)
     idx = inv_areas > 0
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    inv_areas[idx] = 1.0 / inv_areas[idx]
+    inv_areas[idx] = torch.reciprocal(inv_areas[idx])
     inv_areas = inv_areas.view(-1, 1)
 
     return L, inv_areas
@@ -167,7 +160,7 @@ def norm_laplacian(
     v0, v1 = edge_verts[:, 0], edge_verts[:, 1]
 
     # Side lengths of each edge, of shape (E,)
-    w01 = 1.0 / ((v0 - v1).norm(dim=1) + eps)
+    w01 = torch.reciprocal((v0 - v1).norm(dim=1) + eps)
 
     # Construct a sparse matrix by basically doing:
     # L[v0, v1] = w01
@@ -175,8 +168,7 @@ def norm_laplacian(
     e01 = edges.t()  # (2, E)
 
     V = verts.shape[0]
-    # pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
-    L = torch.sparse.FloatTensor(e01, w01, (V, V))
+    L = torch.sparse_coo_tensor(e01, w01, (V, V), dtype=torch.float32)
     L = L + L.t()
 
     return L

pytorch3d/ops/sample_farthest_points.py

@@ -55,6 +55,7 @@ def sample_farthest_points(
     N, P, D = points.shape
     device = points.device
 
+    constant_length = lengths is None
     # Validate inputs
     if lengths is None:
         lengths = torch.full((N,), P, dtype=torch.int64, device=device)
@@ -65,7 +66,9 @@ def sample_farthest_points(
             raise ValueError("A value in lengths was too large.")
 
     # TODO: support providing K as a ratio of the total number of points instead of as an int
+    max_K = -1
     if isinstance(K, int):
+        max_K = K
         K = torch.full((N,), K, dtype=torch.int64, device=device)
     elif isinstance(K, list):
         K = torch.tensor(K, dtype=torch.int64, device=device)
@@ -82,15 +85,19 @@ def sample_farthest_points(
         K = K.to(torch.int64)
 
     # Generate the starting indices for sampling
-    start_idxs = torch.zeros_like(lengths)
     if random_start_point:
-        for n in range(N):
-            # pyre-fixme[6]: For 1st param expected `int` but got `Tensor`.
-            start_idxs[n] = torch.randint(high=lengths[n], size=(1,)).item()
+        if constant_length:
+            start_idxs = torch.randint(high=P, size=(N,), device=device)
+        else:
+            start_idxs = (lengths * torch.rand(lengths.size(), device=device)).to(
+                torch.int64
+            )
+    else:
+        start_idxs = torch.zeros_like(lengths)
 
     with torch.no_grad():
         # pyre-fixme[16]: `pytorch3d_._C` has no attribute `sample_farthest_points`.
-        idx = _C.sample_farthest_points(points, lengths, K, start_idxs)
+        idx = _C.sample_farthest_points(points, lengths, K, start_idxs, max_K)
     sampled_points = masked_gather(points, idx)
 
     return sampled_points, idx

pytorch3d/ops/sample_points_from_meshes.py

@@ -16,9 +16,7 @@ import sys
 from typing import Tuple, Union
 
 import torch
-
 from pytorch3d.ops.mesh_face_areas_normals import mesh_face_areas_normals
-
 from pytorch3d.ops.packed_to_padded import packed_to_padded
 from pytorch3d.renderer.mesh.rasterizer import Fragments as MeshFragments
 

pytorch3d/renderer/__init__.py

@@ -69,7 +69,6 @@ from .mesh import (
     TexturesUV,
     TexturesVertex,
 )
-
 from .points import (
     AlphaCompositor,
     NormWeightedCompositor,

pytorch3d/renderer/camera_conversions.py

@@ -153,12 +153,12 @@ def _pulsar_from_opencv_projection(
     # Check image sizes.
     image_w = image_size_wh[0, 0]
     image_h = image_size_wh[0, 1]
-    assert torch.all(
-        image_size_wh[:, 0] == image_w
-    ), "All images in a batch must have the same width!"
-    assert torch.all(
-        image_size_wh[:, 1] == image_h
-    ), "All images in a batch must have the same height!"
+    assert torch.all(image_size_wh[:, 0] == image_w), (
+        "All images in a batch must have the same width!"
+    )
+    assert torch.all(image_size_wh[:, 1] == image_h), (
+        "All images in a batch must have the same height!"
+    )
     # Focal length.
     fx = camera_matrix[:, 0, 0].unsqueeze(1)
     fy = camera_matrix[:, 1, 1].unsqueeze(1)

pytorch3d/renderer/cameras.py

@@ -629,10 +629,8 @@ class FoVPerspectiveCameras(CamerasBase):
         # so the so the z sign is 1.0.
         z_sign = 1.0
 
-        # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-        K[:, 0, 0] = 2.0 * znear / (max_x - min_x)
-        # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-        K[:, 1, 1] = 2.0 * znear / (max_y - min_y)
+        K[:, 0, 0] = torch.div(2.0 * znear, max_x - min_x)
+        K[:, 1, 1] = torch.div(2.0 * znear, max_y - min_y)
         K[:, 0, 2] = (max_x + min_x) / (max_x - min_x)
         K[:, 1, 2] = (max_y + min_y) / (max_y - min_y)
         K[:, 3, 2] = z_sign * ones
@@ -1178,9 +1176,7 @@ class PerspectiveCameras(CamerasBase):
         xy_inv_depth = torch.cat(
             # pyre-fixme[6]: For 1st argument expected `Union[List[Tensor],
             #  tuple[Tensor, ...]]` but got `Tuple[Tensor, float]`.
-            # pyre-fixme[58]: `/` is not supported for operand types `float` and
-            #  `Tensor`.
-            (xy_depth[..., :2], 1.0 / xy_depth[..., 2:3]),
+            (xy_depth[..., :2], torch.reciprocal(xy_depth[..., 2:3])),
             dim=-1,  # type: ignore
         )
         return unprojection_transform.transform_points(xy_inv_depth)

pytorch3d/renderer/mesh/__init__.py

@@ -12,7 +12,6 @@ from .clip import (
     ClippedFaces,
     convert_clipped_rasterization_to_original_faces,
 )
-
 from .rasterize_meshes import rasterize_meshes
 from .rasterizer import MeshRasterizer, RasterizationSettings
 from .renderer import MeshRenderer, MeshRendererWithFragments

pytorch3d/renderer/mesh/clip.py

@@ -434,13 +434,7 @@ def clip_faces(
     # These will then be filled in for each case.
     ###########################################
     F_clipped = (
-        F
-        # pyre-fixme[58]: `+` is not supported for operand types `int` and
-        #  `Union[bool, float, int]`.
-        + faces_delta_cum[-1].item()
-        # pyre-fixme[58]: `+` is not supported for operand types `int` and
-        #  `Union[bool, float, int]`.
-        + faces_delta[-1].item()
+        F + int(faces_delta_cum[-1].item()) + int(faces_delta[-1].item())
     )  # Total number of faces in the new Meshes
     face_verts_clipped = torch.zeros(
         (F_clipped, 3, 3), dtype=face_verts_unclipped.dtype, device=device

pytorch3d/renderer/mesh/rasterize_meshes.py

@@ -14,7 +14,6 @@ import torch
 from pytorch3d import _C
 
 from ..utils import parse_image_size
-
 from .clip import (
     clip_faces,
     ClipFrustum,

pytorch3d/renderer/mesh/textures.py

@@ -71,9 +71,7 @@ def _list_to_padded_wrapper(
         # pyre-fixme[6]: For 2nd param expected `int` but got `Union[bool, float, int]`.
         x_reshaped.append(y.reshape(-1, D))
     x_padded = list_to_padded(x_reshaped, pad_size=pad_size, pad_value=pad_value)
-    # pyre-fixme[58]: `+` is not supported for operand types `Tuple[int, int]` and
-    #  `Size`.
-    return x_padded.reshape((N, -1) + reshape_dims)
+    return x_padded.reshape((N, -1) + tuple(reshape_dims))
 
 
 def _padded_to_list_wrapper(
@@ -104,9 +102,7 @@ def _padded_to_list_wrapper(
     # pyre-fixme[6]: For 3rd param expected `int` but got `Union[bool, float, int]`.
     x_reshaped = x.reshape(N, M, D)
     x_list = padded_to_list(x_reshaped, split_size=split_size)
-    # pyre-fixme[58]: `+` is not supported for operand types `Tuple[typing.Any]` and
-    #  `Size`.
-    x_list = [xl.reshape((xl.shape[0],) + reshape_dims) for xl in x_list]
+    x_list = [xl.reshape((xl.shape[0],) + tuple(reshape_dims)) for xl in x_list]
     return x_list
 
 
@@ -625,9 +621,7 @@ class TexturesAtlas(TexturesBase):
         of length `k`.
         """
         if len(faces_ids_list) != len(self.atlas_list()):
-            raise IndexError(
-                "faces_ids_list must be of " "the same length as atlas_list."
-            )
+            raise IndexError("faces_ids_list must be of the same length as atlas_list.")
 
         sub_features = []
         for atlas, faces_ids in zip(self.atlas_list(), faces_ids_list):
@@ -1657,7 +1651,7 @@ class TexturesUV(TexturesBase):
             raise NotImplementedError("This function does not support multiple maps.")
         if len(faces_ids_list) != len(self.faces_uvs_padded()):
             raise IndexError(
-                "faces_uvs_padded must be of " "the same length as face_ids_list."
+                "faces_uvs_padded must be of the same length as face_ids_list."
             )
 
         sub_faces_uvs, sub_verts_uvs, sub_maps = [], [], []
@@ -1871,7 +1865,7 @@ class TexturesVertex(TexturesBase):
         """
         if len(vertex_ids_list) != len(self.verts_features_list()):
             raise IndexError(
-                "verts_features_list must be of " "the same length as vertex_ids_list."
+                "verts_features_list must be of the same length as vertex_ids_list."
             )
 
         sub_features = []

pytorch3d/renderer/opengl/opengl_utils.py

@@ -24,7 +24,6 @@ from typing import Any, Dict
 
 os.environ["PYOPENGL_PLATFORM"] = "egl"
 import OpenGL.EGL as egl  # noqa
-
 import pycuda.driver as cuda  # noqa
 from OpenGL._opaque import opaque_pointer_cls  # noqa
 from OpenGL.raw.EGL._errors import EGLError  # noqa

pytorch3d/renderer/opengl/rasterizer_opengl.py

@@ -17,15 +17,12 @@ import numpy as np
 import OpenGL.GL as gl
 import pycuda.gl
 import torch
-
 import torch.nn as nn
-
 from pytorch3d.structures.meshes import Meshes
 
 from ..cameras import FoVOrthographicCameras, FoVPerspectiveCameras
 from ..mesh.rasterizer import Fragments, RasterizationSettings
 from ..utils import parse_image_size
-
 from .opengl_utils import _torch_to_opengl, global_device_context_store
 
 # Shader strings, used below to compile an OpenGL program.

pytorch3d/renderer/points/__init__.py

@@ -9,9 +9,7 @@
 import torch
 
 from .compositor import AlphaCompositor, NormWeightedCompositor
-
 from .pulsar.unified import PulsarPointsRenderer
-
 from .rasterize_points import rasterize_points
 from .rasterizer import PointsRasterizationSettings, PointsRasterizer
 from .renderer import PointsRenderer

pytorch3d/renderer/points/rasterize_points.py

@@ -11,7 +11,6 @@ from typing import List, Optional, Tuple, Union
 import numpy as np
 import torch
 from pytorch3d import _C
-
 from pytorch3d.renderer.mesh.rasterize_meshes import pix_to_non_square_ndc
 
 from ..utils import parse_image_size

pytorch3d/renderer/utils.py

@@ -269,9 +269,7 @@ class TensorProperties(nn.Module):
                         # to have the same shape as the input tensor.
                         new_dims = len(tensor_dims) - len(idx_dims)
                         new_shape = idx_dims + (1,) * new_dims
-                        # pyre-fixme[58]: `+` is not supported for operand types
-                        # `Tuple[int]` and `torch._C.Size`
-                        expand_dims = (-1,) + tensor_dims[1:]
+                        expand_dims = (-1,) + tuple(tensor_dims[1:])
                         _batch_idx = _batch_idx.view(*new_shape)
                         _batch_idx = _batch_idx.expand(*expand_dims)
 

pytorch3d/structures/meshes.py

@@ -531,9 +531,9 @@ class Meshes:
             list of tensors of vertices of shape (V_n, 3).
         """
         if self._verts_list is None:
-            assert (
-                self._verts_padded is not None
-            ), "verts_padded is required to compute verts_list."
+            assert self._verts_padded is not None, (
+                "verts_padded is required to compute verts_list."
+            )
             self._verts_list = struct_utils.padded_to_list(
                 self._verts_padded, self.num_verts_per_mesh().tolist()
             )
@@ -547,9 +547,9 @@ class Meshes:
             list of tensors of faces of shape (F_n, 3).
         """
         if self._faces_list is None:
-            assert (
-                self._faces_padded is not None
-            ), "faces_padded is required to compute faces_list."
+            assert self._faces_padded is not None, (
+                "faces_padded is required to compute faces_list."
+            )
             self._faces_list = struct_utils.padded_to_list(
                 self._faces_padded, self.num_faces_per_mesh().tolist()
             )
@@ -925,9 +925,9 @@ class Meshes:
 
         verts_list = self.verts_list()
         faces_list = self.faces_list()
-        assert (
-            faces_list is not None and verts_list is not None
-        ), "faces_list and verts_list arguments are required"
+        assert faces_list is not None and verts_list is not None, (
+            "faces_list and verts_list arguments are required"
+        )
 
         if self.isempty():
             self._faces_padded = torch.zeros(

pytorch3d/structures/pointclouds.py

@@ -433,9 +433,9 @@ class Pointclouds:
             list of tensors of points of shape (P_n, 3).
         """
         if self._points_list is None:
-            assert (
-                self._points_padded is not None
-            ), "points_padded is required to compute points_list."
+            assert self._points_padded is not None, (
+                "points_padded is required to compute points_list."
+            )
             points_list = []
             for i in range(self._N):
                 points_list.append(

pytorch3d/transforms/rotation_conversions.py

@@ -52,8 +52,7 @@ def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
         Rotation matrices as tensor of shape (..., 3, 3).
     """
     r, i, j, k = torch.unbind(quaternions, -1)
-    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+    two_s = torch.div(2.0, (quaternions * quaternions).sum(-1))
 
     o = torch.stack(
         (
@@ -137,18 +136,18 @@ def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
     # we produce the desired quaternion multiplied by each of r, i, j, k
     quat_by_rijk = torch.stack(
         [
-            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
-            #  `int`.
-            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
-            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
-            #  `int`.
-            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
-            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
-            #  `int`.
-            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
-            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
-            #  `int`.
-            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+            torch.stack(
+                [torch.square(q_abs[..., 0]), m21 - m12, m02 - m20, m10 - m01], dim=-1
+            ),
+            torch.stack(
+                [m21 - m12, torch.square(q_abs[..., 1]), m10 + m01, m02 + m20], dim=-1
+            ),
+            torch.stack(
+                [m02 - m20, m10 + m01, torch.square(q_abs[..., 2]), m12 + m21], dim=-1
+            ),
+            torch.stack(
+                [m10 - m01, m20 + m02, m21 + m12, torch.square(q_abs[..., 3])], dim=-1
+            ),
         ],
         dim=-2,
     )
@@ -160,9 +159,10 @@ def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
 
     # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
     # forall i; we pick the best-conditioned one (with the largest denominator)
-    out = quat_candidates[
-        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :
-    ].reshape(batch_dim + (4,))
+    indices = q_abs.argmax(dim=-1, keepdim=True)
+    expand_dims = list(batch_dim) + [1, 4]
+    gather_indices = indices.unsqueeze(-1).expand(expand_dims)
+    out = torch.gather(quat_candidates, -2, gather_indices).squeeze(-2)
     return standardize_quaternion(out)
 
 
@@ -293,10 +293,11 @@ def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tenso
     tait_bryan = i0 != i2
     if tait_bryan:
         central_angle = torch.asin(
-            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
+            torch.clamp(matrix[..., i0, i2], -1.0, 1.0)
+            * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
         )
     else:
-        central_angle = torch.acos(matrix[..., i0, i0])
+        central_angle = torch.acos(torch.clamp(matrix[..., i0, i0], -1.0, 1.0))
 
     o = (
         _angle_from_tan(

pytorch3d/transforms/transform3d.py

@@ -623,9 +623,7 @@ class Scale(Transform3d):
         Return the inverse of self._matrix.
         """
         xyz = torch.stack([self._matrix[:, i, i] for i in range(4)], dim=1)
-        # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
-        ixyz = 1.0 / xyz
-        # pyre-fixme[6]: For 1st param expected `Tensor` but got `float`.
+        ixyz = torch.reciprocal(xyz)
         imat = torch.diag_embed(ixyz, dim1=1, dim2=2)
         return imat
 

pytorch3d/utils/__init__.py

@@ -12,11 +12,8 @@ from .camera_conversions import (
     pulsar_from_cameras_projection,
     pulsar_from_opencv_projection,
 )
-
 from .checkerboard import checkerboard
-
 from .ico_sphere import ico_sphere
-
 from .torus import torus
 
 

setup.py

@@ -75,6 +75,21 @@ def get_extensions():
         ]
         if os.name != "nt":
             nvcc_args.append("-std=c++17")
+
+        # CUDA 13.0+ compatibility flags for pulsar.
+        # Starting with CUDA 13, __global__ function visibility changed.
+        # See: https://developer.nvidia.com/blog/
+        #      cuda-c-compiler-updates-impacting-elf-visibility-and-linkage/
+        cuda_version = torch.version.cuda
+        if cuda_version is not None:
+            major = int(cuda_version.split(".")[0])
+            if major >= 13:
+                nvcc_args.extend(
+                    [
+                        "--device-entity-has-hidden-visibility=false",
+                        "-static-global-template-stub=false",
+                    ]
+                )
         if cub_home is None:
             prefix = os.environ.get("CONDA_PREFIX", None)
             if prefix is not None and os.path.isdir(prefix + "/include/cub"):
@@ -134,7 +149,7 @@ if os.getenv("PYTORCH3D_NO_NINJA", "0") == "1":
 
     class BuildExtension(torch.utils.cpp_extension.BuildExtension):
         def __init__(self, *args, **kwargs):
-            super().__init__(use_ninja=False, *args, **kwargs)
+            super().__init__(*args, use_ninja=False, **kwargs)
 
 else:
     BuildExtension = torch.utils.cpp_extension.BuildExtension

tests/benchmarks/bm_ball_query_large.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.
+
+from itertools import product
+
+import torch
+from fvcore.common.benchmark import benchmark
+from pytorch3d.ops.ball_query import ball_query
+
+
+def ball_query_square(
+    N: int, P1: int, P2: int, D: int, K: int, radius: float, device: str
+):
+    device = torch.device(device)
+    pts1 = torch.rand(N, P1, D, device=device)
+    pts2 = torch.rand(N, P2, D, device=device)
+    torch.cuda.synchronize()
+
+    def output():
+        ball_query(pts1, pts2, K=K, radius=radius, skip_points_outside_cube=True)
+        torch.cuda.synchronize()
+
+    return output
+
+
+def bm_ball_query() -> None:
+    backends = ["cpu", "cuda:0"]
+
+    kwargs_list = []
+    Ns = [32]
+    P1s = [256]
+    P2s = [2**p for p in range(9, 20, 2)]
+    Ds = [3, 10]
+    Ks = [500]
+    Rs = [0.01, 0.1]
+    test_cases = product(Ns, P1s, P2s, Ds, Ks, Rs, backends)
+    for case in test_cases:
+        N, P1, P2, D, K, R, b = case
+        kwargs_list.append(
+            {"N": N, "P1": P1, "P2": P2, "D": D, "K": K, "radius": R, "device": b}
+        )
+    benchmark(
+        ball_query_square,
+        "BALLQUERY_SQUARE",
+        kwargs_list,
+        num_iters=30,
+        warmup_iters=1,
+    )
+
+
+if __name__ == "__main__":
+    bm_ball_query()