Bias for rescaling by input (#44)

README.md

@@ -11,6 +11,7 @@
 
 - Exponential layer
 - Flexible multi-layer neural network with optional nonlinear last layer
+- Affine rescaling of output by an input
 
 
 ## Development

econ_layers/layers.py

@@ -20,15 +20,20 @@ def forward(self, input):
 
 # rescaling by a specific element of a given input
 class RescaleOutputsByInput(nn.Module):
-    def __init__(self, rescale_index: int = 0):
+    def __init__(self, rescale_index: int = 0, bias=False):
         super().__init__()
         self.rescale_index = rescale_index
-
+        if bias:
+            self.bias =  torch.nn.Parameter(torch.Tensor(1)) # only a scalar here
+            torch.nn.init.zeros_(self.bias)     
+        else:
+            self.bias = 0.0 # register_parameter('bias', None) # necessary?
+
     def forward(self, x, y):
         if x.dim() == 1:
-            return x[self.rescale_index] * y
+            return x[self.rescale_index] * y + self.bias
         else:
-            return x[:, [self.rescale_index]] * y
+            return x[:, [self.rescale_index]] * y + self.bias
 
 
 # assuming 2D data

setup.py

@@ -57,6 +57,6 @@
     test_suite="tests",
     tests_require=test_requirements,
     url="https://github.com/HighDimensionalEconLab/econ_layers",
-    version="0.0.23",
+    version="0.0.24",
     zip_safe=False,
 )

tests/test_input_rescaling.py

@@ -5,8 +5,19 @@
 import pytest
 import torch
 import numpy
-from econ_layers.layers import RescaleOutputsByInput
-
+from econ_layers.layers import (
+    FlexibleSequential,
+    RescaleOutputsByInput,
+    ScalarExponentialRescaling,
+    RescaleOutputsByInput,
+)
+import numpy.testing
+import torch.autograd.gradcheck
+from torch.autograd import Variable
+from tests.helpers import train, test
+
+torch.set_printoptions(16)  # to be able to see what is going on
+torch.manual_seed(0)
 
 def test_input_rescaling():
 
@@ -33,4 +44,15 @@ def test_input_rescaling():
     assert torch.all(torch.isclose(input_mult_0, x[:, [0]] * x[:, [1]]))
     assert torch.all(torch.isclose(input_mult_1, x[:, [0]] * x[:, [1]]))
 
-
+def test_input_rescaling_bias():
+    n_in = 2
+    n_out = 1
+    mod = FlexibleSequential(
+        n_in,
+        n_out,
+        layers=2,
+        hidden_dim=128,
+        OutputRescalingLayer=RescaleOutputsByInput(rescale_index=0, bias = True),
+    ).double()
+    input = (Variable(torch.randn(n_in).double(), requires_grad=True),)
+    assert torch.autograd.gradcheck(mod, input)