Add type error suppressions for upcoming upgrade

Reviewed By: MaggieMoss

Differential Revision: D77242941

fbshipit-source-id: 08b3645f9b4189ebc2da2ed80c9f71abe4d5926d

Author: generatedunixname89002005307016

flowtorch/bijectors/permute.py

@@ -46,6 +46,7 @@ def _inverse(
         if self.permutation is None:
             self.permutation = torch.randperm(y.shape[-1])
 
+        # pyre-fixme[6]: For 3rd argument expected `Tensor` but got `Optional[Tensor]`.
         x = torch.index_select(y, -1, self.inv_permutation)
         ladj = self._log_abs_det_jacobian(x, y, params)
         return x, ladj

flowtorch/distributions/flow.py

@@ -30,6 +30,7 @@ def __init__(
         # TODO: Confirm that the following logic works. Shouldn't it use
         # .domain and .codomain?? Infer shape from constructed self.bijector
         shape = (
+            # pyre-fixme[6]: For 1st argument expected `tuple[int, ...]` but got `Size`.
             self.base_dist.batch_shape
             # pyre-fixme[58]: `+` is not supported for operand types `Size` and `Size`.
             + self.base_dist.event_shape

flowtorch/distributions/neals_funnel.py

@@ -34,7 +34,11 @@ def rsample(
         if not sample_shape:
             sample_shape = torch.Size()
         eps = _standard_normal(
-            (sample_shape[0], 2), dtype=torch.float, device=torch.device("cpu")
+            # pyre-fixme[6]: For 1st argument expected `Sequence[Union[int,
+            #  SymInt]]` but got `Tuple[Union[int, Tensor], int]`.
+            (sample_shape[0], 2),
+            dtype=torch.float,
+            device=torch.device("cpu"),
         )
         z = torch.zeros(eps.shape)
         z[..., 1] = torch.tensor(3.0) * eps[..., 1]

flowtorch/parameters/dense_autoregressive.py

@@ -182,6 +182,7 @@ def _forward(
         # results ~ (batch_shape, param_shapes[0]), ...
         result = tuple(
             # pyre-fixme[58]: `+` is not supported for operand types `Size` and `Size`.
+            # pyre-fixme[6]: For 1st argument expected `tuple[int, ...]` but got `Size`.
             h_slice.view(batch_shape + p_shape)
             for h_slice, p_shape in zip(result, list(self.param_shapes))
         )