Enable Pyrefly in fbcode/pearl

Summary:
Automated migration to enable Pyrefly type checking for `fbcode/pearl`.

- Added `python.set_pyrefly(True)` to PACKAGE file
- Suppressed pre-existing type errors

Pyrefly is Meta's next-generation Python type checker, replacing Pyre.

If you encounter issues, you can revert the PACKAGE change by removing
the `python.set_pyrefly(True)` line.
#pyreupgrade

Differential Revision: D102231580

fbshipit-source-id: c7970bedd0c10f341cb1826cf7d74fa863db706c

Author: generatedunixname89002005307016

pearl/history_summarization_modules/lstm_history_summarization_module.py

@@ -60,6 +60,7 @@ def summarize_history(
             observation.clone().detach().float().view((1, self.observation_dim))
         )
         if action is None:
+            # pyrefly: ignore [bad-assignment]
             action = self.default_action
         assert isinstance(action, torch.Tensor)
         action = action.clone().detach().float().view((1, self.action_dim))
@@ -130,6 +131,7 @@ def compare(self, other: HistorySummarizationModule) -> str:
             differences.append(
                 f"action_dim is different: {self.action_dim} vs {other.action_dim}"
             )
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self.default_action, other.default_action):
             differences.append(
                 f"default_action is different: {self.default_action} vs {other.default_action}"

pearl/history_summarization_modules/stacking_history_summarization_module.py

@@ -40,6 +40,7 @@ def summarize_history(
         self, observation: Observation, action: Action | None
     ) -> torch.Tensor:
         if action is None:
+            # pyrefly: ignore [bad-assignment]
             action = self.default_action
 
         observation = assert_is_tensor_like(observation)
@@ -103,6 +104,7 @@ def compare(self, other: HistorySummarizationModule) -> str:
             differences.append(
                 f"action_dim is different: {self.action_dim} vs {other.action_dim}"
             )
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self.default_action, other.default_action):
             differences.append(
                 f"default_action is different: {self.default_action} vs {other.default_action}"

pearl/history_summarization_modules/transformer_history_summarization_module.py

@@ -39,6 +39,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x: [B, T, d_model]
         """
         T = x.size(1)
+        # pyrefly: ignore [bad-index]
         x = x + self.pe[:T, :].unsqueeze(0)
         return self.dropout(x)
 
@@ -137,6 +138,7 @@ def __init__(
                 d_model=d_model, dropout=dropout
             )
         elif pos_encoding == "sinusoidal":
+            # pyrefly: ignore [bad-assignment]
             self.pos_encoding = SinusoidalPositionalEncoding(
                 d_model=d_model, dropout=dropout
             )
@@ -206,6 +208,7 @@ def summarize_history(
             act = action.clone().detach().float().view(1, self.action_dim)
 
         # Concatenate (action, observation) for current step
+        # pyrefly: ignore [no-matching-overload]
         pair = torch.cat((act, obs), dim=-1)  # [1, input_dim]
         assert pair.shape[-1] == self.history.shape[-1]
 
@@ -279,6 +282,7 @@ def compare(self, other: HistorySummarizationModule) -> str:
             )
 
         # Buffers
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self.default_action, other.default_action):
             differences.append(
                 f"default_action is different: {self.default_action} vs {other.default_action}"
@@ -288,6 +292,7 @@ def compare(self, other: HistorySummarizationModule) -> str:
 
         # Positional encoding (buffer + dropout p)
         if hasattr(self.pos_encoding, "pe") and hasattr(other.pos_encoding, "pe"):
+            # pyrefly: ignore [bad-argument-type]
             if not torch.allclose(self.pos_encoding.pe, other.pos_encoding.pe):
                 differences.append("positional_encoding.pe is different")
         p_self = getattr(self.pos_encoding.dropout, "p", None)

pearl/neural_networks/common/__init__.py

@@ -18,5 +18,6 @@
     "ValueNetwork",
     "CNNValueNetwork",
     "VanillaValueNetwork",
+    # pyrefly: ignore [bad-dunder-all]
     "Epinet",
 ]

pearl/neural_networks/common/epistemic_neural_networks.py

@@ -165,6 +165,7 @@ def __init__(
         self.models: nn.ModuleList = nn.ModuleList(models)
 
         self.params: dict[str, Any]
+        # pyrefly: ignore [bad-override]
         self.buffers: dict[str, Any]
         self.generate_params_buffers()
 

pearl/neural_networks/common/utils.py

@@ -111,11 +111,15 @@ def mlp_block(
             nn.Linear(input_dim_current_layer, output_dim_current_layer)
         )
         if use_layer_norm:
+            # pyrefly: ignore [bad-argument-type]
             single_layers.append(nn.LayerNorm(output_dim_current_layer))
         if dropout_ratio > 0:
+            # pyrefly: ignore [bad-argument-type]
             single_layers.append(nn.Dropout(p=dropout_ratio))
+        # pyrefly: ignore [bad-argument-type]
         single_layers.append(ActivationType(hidden_activation).module())
         if use_batch_norm:
+            # pyrefly: ignore [bad-argument-type]
             single_layers.append(nn.BatchNorm1d(output_dim_current_layer))
         single_layer_model = nn.Sequential(*single_layers)
         if use_skip_connections:
@@ -133,6 +137,7 @@ def mlp_block(
     last_layer = []
     last_layer.append(nn.Linear(dims[-2], dims[-1]))
     if last_activation is not None:
+        # pyrefly: ignore [bad-argument-type]
         last_layer.append(ActivationType(last_activation).module())
     last_layer_model = nn.Sequential(*last_layer)
     if use_skip_connections:
@@ -184,7 +189,9 @@ def conv_block(
         layers.append(conv_layer)
         if use_batch_norm:
             # batch norm should normalize the output of the convolutional layer
+            # pyrefly: ignore [bad-argument-type]
             layers.append(nn.BatchNorm2d(out_channels))
+        # pyrefly: ignore [bad-argument-type]
         layers.append(nn.ReLU())
         # number of input channels to next layer is number of output channels of previous layer:
         input_channels_count = out_channels

pearl/neural_networks/contextual_bandit/linear_regression.py

@@ -90,6 +90,7 @@ def __init__(
     @property
     def A(self) -> torch.Tensor:
         # return A with L2 regularization applied
+        # pyrefly: ignore [no-matching-overload]
         return self._A + self.l2_reg_lambda * torch.eye(
             self._feature_dim + 1, device=self._A.device
         )
@@ -209,8 +210,11 @@ def learn_batch(
             torch.distributed.all_reduce(delta_b)
             torch.distributed.all_reduce(delta_sum_weight)
 
+        # pyrefly: ignore [no-matching-overload]
         self._A += delta_A.to(self._A.device)
+        # pyrefly: ignore [no-matching-overload]
         self._b += delta_b.to(self._b.device)
+        # pyrefly: ignore [no-matching-overload]
         self._sum_weight += delta_sum_weight.to(self._sum_weight.device)
 
         self.calculate_coefs()  # update coefs after updating A and b
@@ -226,7 +230,9 @@ def apply_discounting(self) -> None:
         """
         if self.gamma < 1:
             logger.info(f"Applying discounting at sum_weight={self._sum_weight}")
+            # pyrefly: ignore [bad-argument-type, unsupported-operation]
             self._A *= self.gamma
+            # pyrefly: ignore [bad-argument-type, unsupported-operation]
             self._b *= self.gamma
         # don't dicount sum_weight because it's used to determine when to apply discounting
 
@@ -250,6 +256,7 @@ def calculate_coefs(self) -> None:
         Save inverted A and coefficients in buffers.
         """
         self._inv_A = self.matrix_inv_fallback_pinv(self.A)
+        # pyrefly: ignore [bad-argument-type]
         self._coefs = torch.matmul(self._inv_A, self._b)
 
     def calculate_sigma(self, x: torch.Tensor) -> torch.Tensor:
@@ -300,10 +307,13 @@ def compare(self, other: MuSigmaCBModel) -> str:
                 f"distribution_enabled is different: {self.distribution_enabled} "
                 + f"vs {other.distribution_enabled}"
             )
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self._A, other._A):
             differences.append(f"_A is different: {self._A} vs {other._A}")
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self._b, other._b):
             differences.append(f"_b is different: {self._b} vs {other._b}")
+        # pyrefly: ignore [bad-argument-type]
         if not torch.allclose(self._sum_weight, other._sum_weight):
             differences.append(
                 f"_sum_weight is different: {self._sum_weight} vs {other._sum_weight}"

pearl/neural_networks/sequential_decision_making/actor_networks.py

@@ -577,6 +577,7 @@ def sample_action(
 
         log_prob = normal.log_prob(sample)
         log_prob -= torch.log(
+            # pyrefly: ignore [unsupported-operation]
             self._action_bound * (1 - normalized_action.pow(2)) + epsilon
         )
 
@@ -617,6 +618,7 @@ def get_log_probability(
         unnormalized_action_batch = torch.atanh(normalized_action_batch)
         log_prob = normal.log_prob(unnormalized_action_batch)
         log_prob -= torch.log(
+            # pyrefly: ignore [unsupported-operation]
             self._action_bound * (1 - normalized_action_batch.pow(2)) + epsilon
         )
 

pearl/neural_networks/sequential_decision_making/q_value_networks.py

@@ -294,6 +294,7 @@ def __init__(
         )
         self.register_buffer(
             "_quantile_midpoints",
+            # pyrefly: ignore [bad-index]
             ((self._quantiles[1:] + self._quantiles[:-1]) / 2)
             .unsqueeze(0)
             .unsqueeze(0),
@@ -328,10 +329,12 @@ def get_q_value_distribution(
 
     @property
     def quantiles(self) -> Tensor:
+        # pyrefly: ignore [bad-return]
         return self._quantiles
 
     @property
     def quantile_midpoints(self) -> Tensor:
+        # pyrefly: ignore [bad-return]
         return self._quantile_midpoints
 
     @property

pearl/policy_learners/contextual_bandits/contextual_bandit_base.py

@@ -60,6 +60,7 @@ def learn_batch(self, batch: TransitionBatch) -> dict[str, Any]:
         pass
 
     @abstractmethod
+    # pyrefly: ignore [bad-override]
     def act(
         self,
         subjective_state: SubjectiveState,