Fix __constants__ for some nn modules (#21071)

davidriazati · facebook-github-bot · commit 736bf7b46c04 · 2019-05-29T13:55:53.000-07:00
Summary: A bunch of modules were missing entries for `__constants__` which was making their `__repr__`s not work. Others had `__constants__` that were not necessary since it was provided by some parent class instead. Fixes #20978 ](https://our.intern.facebook.com/intern/diff/15539518/) Pull Request resolved: #21071 Pulled By: driazati Differential Revision: D15539518 fbshipit-source-id: 24bdd1ef41ef636eefd5d2bad4ab2d79646ed4f0
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -14874,16 +14874,22 @@ class TheModule(torch.jit.ScriptModule):
                 def __init__(self):
                     super(TheModule, self).__init__()
                     self.submodule = nn_module(*constructor_args)
+
+            def make_module(script):
+                module = TheModule()
+                # check __repr__
+                str(module)
+                module.define(script)
+                return module
+
             # module cannot be imported / exported
             if module_name in EXCLUDE_MODULE_EXPORT_IMPORT:
                 with self.disableEmitHook():
-                    module = TheModule()
-                    module.define(script)
+                    module = make_module(script)
                     create_script_module.last_graph = module.graph
                     mod = module(*args)
             else:
-                module = TheModule()
-                module.define(script)
+                module = make_module(script)
                 self.assertExportImportModule(module, tensors)
                 create_script_module.last_graph = module.graph
                 mod = module(*args)
diff --git a/torch/nn/modules/activation.py b/torch/nn/modules/activation.py
@@ -829,6 +829,7 @@ class PReLU(Module):
         >>> input = torch.randn(2)
         >>> output = m(input)
     """
+    __constants__ = ['num_parameters']
 
     def __init__(self, num_parameters=1, init=0.25):
         self.num_parameters = num_parameters
diff --git a/torch/nn/modules/batchnorm.py b/torch/nn/modules/batchnorm.py
@@ -15,7 +15,8 @@
 class _BatchNorm(Module):
     _version = 2
     __constants__ = ['track_running_stats', 'momentum', 'eps', 'weight', 'bias',
-                     'running_mean', 'running_var', 'num_batches_tracked']
+                     'running_mean', 'running_var', 'num_batches_tracked',
+                     'num_features', 'affine']
 
     def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
                  track_running_stats=True):
diff --git a/torch/nn/modules/conv.py b/torch/nn/modules/conv.py
@@ -12,7 +12,9 @@
 @weak_module
 class _ConvNd(Module):
 
-    __constants__ = ['stride', 'padding', 'dilation', 'groups', 'bias', 'padding_mode']
+    __constants__ = ['stride', 'padding', 'dilation', 'groups', 'bias',
+                     'padding_mode', 'output_padding', 'in_channels',
+                     'out_channels', 'kernel_size']
 
     def __init__(self, in_channels, out_channels, kernel_size, stride,
                  padding, dilation, transposed, output_padding,
@@ -478,10 +480,6 @@ def forward(self, input):
 
 @weak_module
 class _ConvTransposeMixin(object):
-    __constants__ = ['stride', 'padding', 'kernel_size', 'dim_size',
-                     'output_padding', 'groups', 'dilation', 'transposed',
-                     'bias', 'padding_mode']
-
     @weak_script_method
     def forward(self, input, output_size=None):
         # type(Tensor, Optional[List[int]]) -> Tensor
diff --git a/torch/nn/modules/instancenorm.py b/torch/nn/modules/instancenorm.py
@@ -4,9 +4,6 @@
 
 
 class _InstanceNorm(_BatchNorm):
-    __constants__ = ['running_mean', 'running_var', 'weight', 'bias',
-                     'track_running_stats', 'momentum', 'eps']
-
     def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=False,
                  track_running_stats=False):
         super(_InstanceNorm, self).__init__(
diff --git a/torch/nn/modules/normalization.py b/torch/nn/modules/normalization.py
@@ -129,7 +129,7 @@ class LayerNorm(Module):
 
     .. _`Layer Normalization`: https://arxiv.org/abs/1607.06450
     """
-    __constants__ = ['normalized_shape', 'weight', 'bias', 'eps']
+    __constants__ = ['normalized_shape', 'weight', 'bias', 'eps', 'elementwise_affine']
 
     def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
         super(LayerNorm, self).__init__()
diff --git a/torch/nn/modules/sparse.py b/torch/nn/modules/sparse.py
@@ -252,7 +252,7 @@ class EmbeddingBag(Module):
         tensor([[-0.8861, -5.4350, -0.0523],
                 [ 1.1306, -2.5798, -1.0044]])
     """
-    __constants__ = ['num_embeddings, embedding_dim', 'max_norm', 'norm_type',
+    __constants__ = ['num_embeddings', 'embedding_dim', 'max_norm', 'norm_type',
                      'scale_grad_by_freq', 'mode', 'sparse']
 
     def __init__(self, num_embeddings, embedding_dim,
