Move all torch.nn.modules type annotations inline (pytorch#38211)

Summary:
Pull Request resolved: pytorch#38211

Just because the annotations are inline doesn't mean the files type
check; most of the newly annotated files have type errors and I
added exclusions for them in mypy.ini.  The payoff of moving
all of these modules inline is I can delete the relevant code
generation logic for the pyi files (which was added ignore
annotations that weren't actually relevant anymore.)

For the most part the translation was completely mechanical, but there
were two hairy issues.  First, I needed to work around a Python 3.6 and
earlier bug where Generic has a nontrivial metaclass.  This fix is in
torch/jit/__init__.py.  Second, module.py, we need to apply the same
fix for avoiding contravariance checks that the pyi file used to have;
this is done by declaring forward as a variable (rather than a
function), which appears to be sufficient enough to get mypy to not
contravariantly check input arguments.

Because we aren't actually typechecking these modules in most
cases, it is inevitable that some of these type annotations are wrong.
I slavishly copied the old annotations from the pyi files unless there
was an obvious correction I could make.  These annotations will probably
need fixing up later.

Signed-off-by: Edward Z. Yang <[email protected]>

Test Plan: Imported from OSS

Differential Revision: D21497397

Pulled By: ezyang

fbshipit-source-id: 2b08bacc152c48f074e7edc4ee5dce1b77d83702

Author: ezyang

aten/src/ATen/native/quantized/cpu/qbatch_norm.cpp

@@ -42,14 +42,19 @@ void compute_fused_params(
 template <bool ReluFused>
 Tensor q_batch_norm2d_impl(
     Tensor qx,
-    Tensor weight,
-    Tensor bias,
+    c10::optional<Tensor> mb_weight,
+    c10::optional<Tensor> mb_bias,
     Tensor mean,
     Tensor var,
     double eps,
     double output_scale,
     int64_t output_zero_point) {
 
+  TORCH_CHECK(mb_weight.has_value(), "Weight must be provided");
+  TORCH_CHECK(mb_bias.has_value(), "Bias must be provided");
+  const auto& weight = *mb_weight;
+  const auto& bias = *mb_bias;
+
   if (qx.numel() == 0) {
     auto out = qx.clone();
     return out;
@@ -131,14 +136,20 @@ Tensor q_batch_norm2d_impl(
 template <bool ReluFused>
 Tensor q_batch_norm3d_impl(
     Tensor qx,
-    Tensor weight,
-    Tensor bias,
+    c10::optional<Tensor> mb_weight,
+    c10::optional<Tensor> mb_bias,
     Tensor mean,
     Tensor var,
     double eps,
     double output_scale,
     int64_t output_zero_point) {
 
+  TORCH_CHECK(mb_weight.has_value(), "Weight must be provided")
+  TORCH_CHECK(mb_bias.has_value(), "Bias must be provided")
+
+  const auto& weight = *mb_weight;
+  const auto& bias = *mb_bias;
+
   if (qx.numel() == 0) {
     auto out = qx.clone();
     return out;
@@ -231,8 +242,12 @@ Tensor quantized_batch_norm(
     double output_scale,
     int64_t output_zero_point) {
   Tensor qy;
+  // TODO: this should arguably support 3d as well
   qy = q_batch_norm2d_impl<false>(
-      qx, weight, bias, mean, var, eps, output_scale, output_zero_point);
+      qx,
+      weight.defined() ? c10::make_optional(weight) : c10::nullopt,
+      bias.defined() ? c10::make_optional(bias) : c10::nullopt,
+      mean, var, eps, output_scale, output_zero_point);
   return qy;
 }
 

aten/src/ATen/native/quantized/cpu/qnormalization.cpp

@@ -123,33 +123,43 @@ TORCH_LIBRARY_IMPL(quantized, QuantizedCPU, m) {
   m.impl("layer_norm", [](
     Tensor input,
     std::vector<int64_t> normalized_shape,  // because IntArrayRef doesn't work
-    Tensor weight /* optional */,
-    Tensor bias /* optional */,
+    c10::optional<Tensor> weight,
+    c10::optional<Tensor> bias,
     double eps,
     double output_scale,
     int64_t output_zero_point) {
-      return quantized_layer_norm_impl(input, normalized_shape, weight, bias, eps, output_scale, output_zero_point);
+      return quantized_layer_norm_impl(
+          input, normalized_shape,
+          weight.has_value() ? *weight : Tensor(),
+          bias.has_value() ? *bias : Tensor(),
+          eps, output_scale, output_zero_point);
   });
   m.impl("group_norm", [](
       Tensor qx,
       int64_t num_groups,
-      Tensor weight,
-      Tensor bias,
+      c10::optional<Tensor> weight,
+      c10::optional<Tensor> bias,
       double eps,
       double output_scale,
       int64_t output_zero_point) {
     return quantized_group_norm_impl(
-        qx, num_groups, weight, bias, eps, output_scale, output_zero_point);
+        qx, num_groups,
+        weight.has_value() ? *weight : Tensor(),
+        bias.has_value() ? *bias : Tensor(),
+        eps, output_scale, output_zero_point);
   });
   m.impl("instance_norm", [](
       Tensor qx,
-      Tensor weight,
-      Tensor bias,
+      c10::optional<Tensor> weight,
+      c10::optional<Tensor> bias,
       double eps,
       double output_scale,
       int64_t output_zero_point) {
     return quantized_instance_norm_impl(
-        qx, weight, bias, eps, output_scale, output_zero_point);
+        qx,
+        weight.has_value() ? *weight : Tensor(),
+        bias.has_value() ? *bias : Tensor(),
+        eps, output_scale, output_zero_point);
   });
 }