WIP

econml/iv/nnet/_deepiv.py

@@ -94,11 +94,14 @@ def mog_loss_model(n_components, d_t):
     # Use logsumexp for numeric stability:
     # LL = C - log(sum(exp(-d2/(2*sig^2) + log(pi_i/sig^d))))
     def make_logloss(d2, sig, pi):
+        #values = pi / K.pow(sig, d_t) * K.exp(-d2 / (2 * K.square(sig)))
+        # return -K.log(K.sum(values, axis=-1))
+
         # logsumexp doesn't exist in keras 2.4; simulate it
         values = - d2 / (2 * K.square(sig)) + K.log(pi / K.pow(sig, d_t))
         # logsumexp(a,b,c) = log(exp(a)+exp(b)+exp(c)) = log((exp(a-k)+exp(b-k)+exp(c-k))*exp(k))
         # = log((exp(a-k)+exp(b-k)+exp(c-k))) + k
-        mx = K.max(values, axis=-1)
+        mx = K.stop_gradient(K.max(values, axis=-1))
         return -K.log(K.sum(K.exp(values - L.Reshape((-1, 1))(mx)), axis=-1)) - mx
 
     ll = L.Lambda(lambda dsp: make_logloss(*dsp), output_shape=(1,))([d2, sig, pi])
@@ -350,7 +353,7 @@ def fit(self, Y, T, X, Z, *, inference=None):
 
         ll = mog_loss_model(n_components, d_t)([pi, mu, sig, t_in])
 
-        model = Model([z_in, x_in, t_in], [ll])
+        model = Model([z_in, x_in, t_in], [])
         model.add_loss(L.Lambda(K.mean)(ll))
         model.compile(self._optimizer)
         # TODO: do we need to give the user more control over other arguments to fit?
@@ -365,7 +368,7 @@ def fit(self, Y, T, X, Z, *, inference=None):
                                  self._n_samples, self._use_upper_bound_loss, self._n_gradient_samples)
 
         rl = lm([z_in, x_in, y_in])
-        response_model = Model([z_in, x_in, y_in], [rl])
+        response_model = Model([z_in, x_in, y_in], [])
         response_model.add_loss(L.Lambda(K.mean)(rl))
         response_model.compile(self._optimizer)
         # TODO: do we need to give the user more control over other arguments to fit?

econml/tests/test_deepiv.py

@@ -10,6 +10,7 @@
 
 import keras
 import keras.backend as K
+import tensorflow
 
 import pytest
 
@@ -34,7 +35,7 @@ def test_stop_grad(self):
         model.compile('nadam')
         model.fit([np.array([[1]]), np.array([[2]]), np.array([[0]])])
 
-    def test_mog_loss(self):
+    def test_aaaamog_loss(self):
         inputs = [keras.layers.Input(shape=s) for s in [(3,), (3, 2), (3,), (2,)]]
         ll_model = keras.engine.Model(inputs, mog_loss_model(3, 2)(inputs))
 
@@ -489,7 +490,7 @@ def datafunction(n, s, images=False, test=False):
         print("losses: {}".format(losses))
 
     @pytest.mark.slow
-    def test_mog_models(self):
+    def test_aamog_models(self):
         d = 2
         n = 5
 
@@ -518,9 +519,10 @@ def norm(lr):
         # pi,mu,sig = MixtureOfGaussians(n, d)(x_network)
         # ll = MixtureOfGaussiansLogLoss(n, d)([pi,mu,sig,t_input])
         model = keras.engine.Model([x_input, t_input], [ll])
-        model.add_loss(K.mean(ll))
+        model.add_loss(ll)
         model.compile('nadam')
-        model.fit([x, t], epochs=5)
+        model.fit([x, t], np.zeros((5000, 1)), epochs=5)
+        llm = model.predict([x, t])
 
         # For some reason this doesn't work at all when run against the CNTK backend...
         # model.compile('nadam', loss=lambda _,l:l)
@@ -542,6 +544,7 @@ def norm(lr):
         t = 10 * np.sin(theta) + np.random.normal(size=(5000, d))
         pi, mu, sig = model2.predict([x])
         sampled_t = model3.predict([x])
+        model.summary()
         llm = model.predict([x, t])
 
         pi_o = np.tile([[0.25, 0.25, 0.25, 0.25, 0]], (5000, 1))