version incompatibility

Installing latest econml throws an error in an attempt to use DeepIV.

installed : keras=='2.15.0', tensorflow=='2.15.0', econml== '0.15.1' in python 3.11.12

Attempted to replicate: https://github.com/py-why/EconML/blob/main/notebooks/Deep%20IV%20Examples.ipynb.

for example:

from econml.iv.nnet import DeepIV
import keras

from econml.iv.nnet import DeepIV
import keras
import numpy as np
import matplotlib.pyplot as plt

keras.__version__
n = 5000

# Initialize exogenous variables; normal errors, uniformly distributed covariates and instruments
e = np.random.normal(size=(n,))
x = np.random.uniform(low=0.0, high=10.0, size=(n,))
z = np.random.uniform(low=0.0, high=10.0, size=(n,))

# Outcome equation
y = t*t / 10 - x*t / 10 + e
# Initialize treatment variable
t = np.sqrt((x+2) * z) + e

treatment_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(64, activation='relu'),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(32, activation='relu'),
                                    keras.layers.Dropout(0.17)])
response_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(64, activation='relu'),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(32, activation='relu'),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(1)])
est = DeepIV(n_components=10, # Number of gaussians in the mixture density networks)
             m=lambda z, x: treatment_model(keras.layers.concatenate([z, x])), # Treatment model
             h=lambda t, x: response_model(keras.layers.concatenate([t, x])), # Response model
             n_samples=1 # Number of samples used to estimate the response
             )
est.fit(y, t, X=x, Z=z) # Z -> instrumental variables
treatment_effects = est.effect(X_test)

AttributeError Traceback (most recent call last)
Cell In[3], line 34
22 response_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
23 keras.layers.Dropout(0.17),
24 keras.layers.Dense(64, activation='relu'),
(...) 27 keras.layers.Dropout(0.17),
28 keras.layers.Dense(1)])
29 est = DeepIV(n_components=10, # Number of gaussians in the mixture density networks)
30 m=lambda z, x: treatment_model(keras.layers.concatenate([z, x])), # Treatment model
31 h=lambda t, x: response_model(keras.layers.concatenate([t, x])), # Response model
32 n_samples=1 # Number of samples used to estimate the response
33 )
---> 34 est.fit(y, t, X=x, Z=z) # Z -> instrumental variables
35 treatment_effects = est.effect(X_test)

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/econml/_cate_estimator.py:131, in BaseCateEstimator._wrap_fit..call(self, Y, T, inference, *args, **kwargs)
129 inference.prefit(self, Y, T, *args, **kwargs)
130 # call the wrapped fit method
--> 131 m(self, Y, T, *args, **kwargs)
132 self._postfit(Y, T, *args, **kwargs)
133 if inference is not None:
134 # NOTE: we call inference fit after calling the main fit method

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/econml/iv/nnet/_deepiv.py:345, in DeepIV.fit(self, Y, T, X, Z, inference)
341 d_n = K.int_shape(treatment_network)[-1]
343 pi, mu, sig = mog_model(n_components, d_n, d_t)([treatment_network])
--> 345 ll = mog_loss_model(n_components, d_t)([pi, mu, sig, t_in])
347 model = Model([z_in, x_in, t_in], [ll])
348 model.add_loss(L.Lambda(K.mean)(ll))

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/econml/iv/nnet/_deepiv.py:100, in mog_loss_model(n_components, d_t)
97 def make_logloss(d2, sig, pi):
98 return -K.logsumexp(-d2 / (2 * K.square(sig)) + K.log(pi / K.pow(sig, d_t)), axis=-1)
--> 100 ll = L.Lambda(lambda dsp: make_logloss(*dsp), output_shape=(1,))([d2, sig, pi])
102 m = Model([pi, mu, sig, t], [ll])
103 return m

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py:70, in filter_traceback..error_handler(*args, **kwargs)
67 filtered_tb = _process_traceback_frames(e.traceback)
68 # To get the full stack trace, call:
69 # tf.debugging.disable_traceback_filtering()
---> 70 raise e.with_traceback(filtered_tb) from None
71 finally:
72 del filtered_tb

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/econml/iv/nnet/_deepiv.py:100, in mog_loss_model..(dsp)
97 def make_logloss(d2, sig, pi):
98 return -K.logsumexp(-d2 / (2 * K.square(sig)) + K.log(pi / K.pow(sig, d_t)), axis=-1)
--> 100 ll = L.Lambda(lambda dsp: make_logloss(*dsp), output_shape=(1,))([d2, sig, pi])
102 m = Model([pi, mu, sig, t], [ll])
103 return m

File /opt/homebrew/Caskroom/miniforge/base/envs/base-econml/lib/python3.11/site-packages/econml/iv/nnet/_deepiv.py:98, in mog_loss_model..make_logloss(d2, sig, pi)
97 def make_logloss(d2, sig, pi):
---> 98 return -K.logsumexp(-d2 / (2 * K.square(sig)) + K.log(pi / K.pow(sig, d_t)), axis=-1)

AttributeError: Exception encountered when calling layer "lambda_2" (type Lambda).

module 'keras.backend' has no attribute 'logsumexp'

Call arguments received by layer "lambda_2" (type Lambda):
• inputs=['tf.Tensor(shape=(None, 10), dtype=float32)', 'tf.Tensor(shape=(None, 10), dtype=float32)', 'tf.Tensor(shape=(None, 10), dtype=float32)']
• mask=None
• training=None

Reproduce the failure in the Deep IV Examples notebook using the reported Python, Keras, TensorFlow, and EconML versions. Start with econml/iv/nnet/_deepiv.py, especially mog_loss_model and the failing logsumexp call. Done means the example's est.fit(y, t, X=x, Z=z) completes without the reported AttributeError.