Prepare for "Fix type-safety of torch.nn.Module instances": wave 2

Summary: See D52890934

Reviewed By: malfet, r-barnes

Differential Revision: D66245100

fbshipit-source-id: 019058106ac7eaacf29c1c55912922ea55894d23

Author: ezyang

projects/implicitron_trainer/impl/optimizer_factory.py

@@ -123,6 +123,7 @@ def __call__(
         """
         # Get the parameters to optimize
         if hasattr(model, "_get_param_groups"):  # use the model function
+            # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
             p_groups = model._get_param_groups(self.lr, wd=self.weight_decay)
         else:
             p_groups = [

projects/implicitron_trainer/impl/training_loop.py

@@ -395,6 +395,7 @@ def _training_or_validation_epoch(
             ):
                 prefix = f"e{stats.epoch}_it{stats.it[trainmode]}"
                 if hasattr(model, "visualize"):
+                    # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
                     model.visualize(
                         viz,
                         visdom_env_imgs,

pytorch3d/implicitron/dataset/utils.py

@@ -329,6 +329,7 @@ def adjust_camera_to_bbox_crop_(
 
     focal_length_px, principal_point_px = _convert_ndc_to_pixels(
         camera.focal_length[0],
+        # pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
         camera.principal_point[0],
         image_size_wh,
     )
@@ -341,6 +342,7 @@ def adjust_camera_to_bbox_crop_(
     )
 
     camera.focal_length = focal_length[None]
+    # pyre-fixme[16]: `PerspectiveCameras` has no attribute `principal_point`.
     camera.principal_point = principal_point_cropped[None]
 
 
@@ -352,6 +354,7 @@ def adjust_camera_to_image_scale_(
 ) -> PerspectiveCameras:
     focal_length_px, principal_point_px = _convert_ndc_to_pixels(
         camera.focal_length[0],
+        # pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
         camera.principal_point[0],
         original_size_wh,
     )
@@ -368,6 +371,7 @@ def adjust_camera_to_image_scale_(
         image_size_wh_output,
     )
     camera.focal_length = focal_length_scaled[None]
+    # pyre-fixme[16]: `PerspectiveCameras` has no attribute `principal_point`.
     camera.principal_point = principal_point_scaled[None]
 
 

pytorch3d/implicitron/models/feature_extractor/resnet_feature_extractor.py

@@ -142,9 +142,15 @@ def _get_resnet_stage_feature_name(self, stage) -> str:
         return f"res_layer_{stage + 1}"
 
     def _resnet_normalize_image(self, img: torch.Tensor) -> torch.Tensor:
+        # pyre-fixme[58]: `-` is not supported for operand types `Tensor` and
+        #  `Union[Tensor, Module]`.
+        # pyre-fixme[58]: `/` is not supported for operand types `Tensor` and
+        #  `Union[Tensor, Module]`.
         return (img - self._resnet_mean) / self._resnet_std
 
     def get_feat_dims(self) -> int:
+        # pyre-fixme[29]: `Union[(self: TensorBase) -> Tensor, Tensor, Module]` is
+        #  not a function.
         return sum(self._feat_dim.values())
 
     def forward(
@@ -183,7 +189,12 @@ def forward(
             else:
                 imgs_normed = imgs_resized
             #  is not a function.
+            # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
             feats = self.stem(imgs_normed)
+            # pyre-fixme[6]: For 1st argument expected `Iterable[_T1]` but got
+            #  `Union[Tensor, Module]`.
+            # pyre-fixme[6]: For 2nd argument expected `Iterable[_T2]` but got
+            #  `Union[Tensor, Module]`.
             for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                 feats = layer(feats)
                 # just a sanity check below

pytorch3d/implicitron/models/generic_model.py

@@ -478,6 +478,8 @@ def curried_viewpooler(pts):
             )
         custom_args["global_code"] = global_code
 
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a
+        #  function.
         for func in self._implicit_functions:
             func.bind_args(**custom_args)
 
@@ -500,6 +502,8 @@ def curried_viewpooler(pts):
         # Unbind the custom arguments to prevent pytorch from storing
         # large buffers of intermediate results due to points in the
         # bound arguments.
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a
+        #  function.
         for func in self._implicit_functions:
             func.unbind_args()
 

pytorch3d/implicitron/models/global_encoder/autodecoder.py

@@ -71,6 +71,7 @@ def _build_key_map(
         return key_map
 
     def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:
+        # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute `weight`.
         return (self._autodecoder_codes.weight**2).mean()
 
     def get_encoding_dim(self) -> int:
@@ -95,13 +96,15 @@ def forward(self, x: Union[torch.LongTensor, List[str]]) -> Optional[torch.Tenso
                 # pyre-fixme[9]: x has type `Union[List[str], LongTensor]`; used as
                 #  `Tensor`.
                 x = torch.tensor(
+                    # pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, ...
                     [self._key_map[elem] for elem in x],
                     dtype=torch.long,
                     device=next(self.parameters()).device,
                 )
             except StopIteration:
                 raise ValueError("Not enough n_instances in the autodecoder") from None
 
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         return self._autodecoder_codes(x)
 
     def _load_key_map_hook(

pytorch3d/implicitron/models/global_encoder/global_encoder.py

@@ -122,6 +122,7 @@ def forward(
         if frame_timestamp.shape[-1] != 1:
             raise ValueError("Frame timestamp's last dimensions should be one.")
         time = frame_timestamp / self.time_divisor
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         return self._harmonic_embedding(time)
 
     def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:

pytorch3d/implicitron/models/implicit_function/decoding_functions.py

@@ -232,9 +232,14 @@ def forward(self, x: torch.Tensor, z: Optional[torch.Tensor] = None):
             # if the skip tensor is None, we use `x` instead.
             z = x
         skipi = 0
+        # pyre-fixme[6]: For 1st argument expected `Iterable[_T]` but got
+        #  `Union[Tensor, Module]`.
         for li, layer in enumerate(self.mlp):
+            # pyre-fixme[58]: `in` is not supported for right operand type
+            #  `Union[Tensor, Module]`.
             if li in self._input_skips:
                 if self._skip_affine_trans:
+                    # pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, ...
                     y = self._apply_affine_layer(self.skip_affines[skipi], y, z)
                 else:
                     y = torch.cat((y, z), dim=-1)

pytorch3d/implicitron/models/implicit_function/idr_feature_field.py

@@ -141,11 +141,16 @@ def forward(
             self.embed_fn is None and fun_viewpool is None and global_code is None
         ):
             return torch.tensor(
-                [], device=rays_points_world.device, dtype=rays_points_world.dtype
+                [],
+                device=rays_points_world.device,
+                dtype=rays_points_world.dtype,
+                # pyre-fixme[6]: For 2nd argument expected `Union[int, SymInt]` but got
+                #  `Union[Module, Tensor]`.
             ).view(0, self.out_dim)
 
         embeddings = []
         if self.embed_fn is not None:
+            # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             embeddings.append(self.embed_fn(rays_points_world))
 
         if fun_viewpool is not None:
@@ -164,13 +169,19 @@ def forward(
 
         embedding = torch.cat(embeddings, dim=-1)
         x = embedding
+        # pyre-fixme[29]: `Union[(self: TensorBase, other: Union[bool, complex,
+        #  float, int, Tensor]) -> Tensor, Module, Tensor]` is not a function.
         for layer_idx in range(self.num_layers - 1):
             if layer_idx in self.skip_in:
                 x = torch.cat([x, embedding], dim=-1) / 2**0.5
 
+            # pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[An...
             x = self.linear_layers[layer_idx](x)
 
+            # pyre-fixme[29]: `Union[(self: TensorBase, other: Union[bool, complex,
+            #  float, int, Tensor]) -> Tensor, Module, Tensor]` is not a function.
             if layer_idx < self.num_layers - 2:
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
                 x = self.softplus(x)
 
         return x

pytorch3d/implicitron/models/implicit_function/neural_radiance_field.py

@@ -123,8 +123,10 @@ def _get_colors(self, features: torch.Tensor, rays_directions: torch.Tensor):
         # Normalize the ray_directions to unit l2 norm.
         rays_directions_normed = torch.nn.functional.normalize(rays_directions, dim=-1)
         # Obtain the harmonic embedding of the normalized ray directions.
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         rays_embedding = self.harmonic_embedding_dir(rays_directions_normed)
 
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         return self.color_layer((self.intermediate_linear(features), rays_embedding))
 
     @staticmethod
@@ -195,6 +197,8 @@ def forward(
         embeds = create_embeddings_for_implicit_function(
             xyz_world=rays_points_world,
             #  for 2nd param but got `Union[None, torch.Tensor, torch.nn.Module]`.
+            # pyre-fixme[6]: For 2nd argument expected `Optional[(...) -> Any]` but
+            #  got `Union[None, Tensor, Module]`.
             xyz_embedding_function=(
                 self.harmonic_embedding_xyz if self.input_xyz else None
             ),
@@ -206,19 +210,23 @@ def forward(
         )
 
         # embeds.shape = [minibatch x n_src x n_rays x n_pts x self.n_harmonic_functions*6+3]
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         features = self.xyz_encoder(embeds)
         # features.shape = [minibatch x ... x self.n_hidden_neurons_xyz]
         # NNs operate on the flattenned rays; reshaping to the correct spatial size
         # TODO: maybe make the transformer work on non-flattened tensors to avoid this reshape
         features = features.reshape(*rays_points_world.shape[:-1], -1)
 
+        # pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
         raw_densities = self.density_layer(features)
         # raw_densities.shape = [minibatch x ... x 1] in [0-1]
 
         if self.xyz_ray_dir_in_camera_coords:
             if camera is None:
                 raise ValueError("Camera must be given if xyz_ray_dir_in_camera_coords")
 
+            # pyre-fixme[58]: `@` is not supported for operand types `Tensor` and
+            #  `Union[Tensor, Module]`.
             directions = ray_bundle.directions @ camera.R
         else:
             directions = ray_bundle.directions