From 518a190dd3045aba850c4287f654a33b16184e77 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 23 Aug 2022 10:59:23 -0700
Subject: [PATCH] Bias for rescaling by input (#44)

---
 README.md                     |  1 +
 econ_layers/layers.py         | 13 +++++++++----
 setup.py                      |  2 +-
 tests/test_input_rescaling.py | 28 +++++++++++++++++++++++++---
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index d899fcb..5c4a16d 100644
--- a/README.md
+++ b/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
diff --git a/econ_layers/layers.py b/econ_layers/layers.py
index b712533..00ba934 100644
--- a/econ_layers/layers.py
+++ b/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
diff --git a/setup.py b/setup.py
index df041c4..11a6421 100644
--- a/setup.py
+++ b/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,
 )
diff --git a/tests/test_input_rescaling.py b/tests/test_input_rescaling.py
index f384451..6c8e12a 100644
--- a/tests/test_input_rescaling.py
+++ b/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)