AttributeError when fitting example DeepIV model

I'm trying to train a DeepIV model as per the example in the documentation. I have tried doing so using native keras, as well as via tensorflow. In both cases, I encounter an AttributeError, although with slightly different stacktraces.

Given a choice, I would prefer to use tensorflow rather than native keras.

Versions

• Keras: 2.3.1 (using tensorflow backend)
• TensorFlow: 2.3.0
• Econml: 0.10.0

Using keras

Code to replicate the issue:

import keras
from econml.iv.nnet import DeepIV

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)

Stack trace:

AttributeError                            Traceback (most recent call last)
<ipython-input-5-263e9f75519e> in <module>
      7                                     keras.layers.Dropout(0.17),
      8                                     keras.layers.Dense(32, activation='relu'),
----> 9                                     keras.layers.Dropout(0.17)])
     10 response_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
     11                                   keras.layers.Dropout(0.17),

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/engine/sequential.py in __init__(self, layers, name)
     92         if layers:
     93             for layer in layers:
---> 94                 self.add(layer)
     95 
     96     @property

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/engine/sequential.py in add(self, layer)
    164                     # and create the node connecting the current layer
    165                     # to the input layer we just created.
--> 166                     layer(x)
    167                     set_inputs = True
    168             else:

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py in symbolic_fn_wrapper(*args, **kwargs)
     73         if _SYMBOLIC_SCOPE.value:
     74             with get_graph().as_default():
---> 75                 return func(*args, **kwargs)
     76         else:
     77             return func(*args, **kwargs)

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, inputs, **kwargs)
    444                 # Raise exceptions in case the input is not compatible
    445                 # with the input_spec specified in the layer constructor.
--> 446                 self.assert_input_compatibility(inputs)
    447 
    448                 # Collect input shapes to build layer.

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/engine/base_layer.py in assert_input_compatibility(self, inputs)
    308         for x in inputs:
    309             try:
--> 310                 K.is_keras_tensor(x)
    311             except ValueError:
    312                 raise ValueError('Layer ' + self.name + ' was called with '

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py in is_keras_tensor(x)
    693     ```
    694     """
--> 695     if not is_tensor(x):
    696         raise ValueError('Unexpectedly found an instance of type `' +
    697                          str(type(x)) + '`. '

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py in is_tensor(x)
    701 
    702 def is_tensor(x):
--> 703     return isinstance(x, tf_ops._TensorLike) or tf_ops.is_dense_tensor_like(x)
    704 
    705 

AttributeError: module 'tensorflow.python.framework.ops' has no attribute '_TensorLike'

Using TensorFlow

Code to replicate the issue:

import numpy as np
import tensorflow as tf
from econml.iv.nnet import DeepIV

n = 2000
epochs = 2
e = np.random.uniform(low=-0.5, high=0.5, size=(n, 1))
z = np.random.uniform(size=(n, 1))
x = np.random.uniform(size=(n, 1)) + e
p = x + z * e + np.random.uniform(size=(n, 1))
y = p * x + e

treatment_model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(10, activation="relu", input_shape=(2,)),
        tf.keras.layers.Dense(10, activation="relu"),
        tf.keras.layers.Dense(10, activation="relu"),
    ]
)

hmodel = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(10, activation="relu", input_shape=(2,)),
        tf.keras.layers.Dense(10, activation="relu"),
        tf.keras.layers.Dense(1),
    ]
)

deepIv = DeepIV(
    n_components=10,
    m=lambda z, x: treatment_model(tf.keras.layers.concatenate([z, x])),
    h=lambda t, x: hmodel(tf.keras.layers.concatenate([t, x])),
    n_samples=2,
)
deepIv.fit(y, p, X=x, Z=z)

Stack trace:

AttributeError                            Traceback (most recent call last)
<ipython-input-9-a7333673df03> in <module>
     21     n_samples=2
     22 )
---> 23 deepIv.fit(y, p, X=x, Z=z)

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/econml/utilities.py in m(*args, **kwargs)
   1260             if wrong_args:
   1261                 warn(message, category, stacklevel=2)
-> 1262             return to_wrap(*args, **kwargs)
   1263         return m
   1264     return decorator

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/econml/_cate_estimator.py in call(self, Y, T, inference, *args, **kwargs)
    128                 inference.prefit(self, Y, T, *args, **kwargs)
    129             # call the wrapped fit method
--> 130             m(self, Y, T, *args, **kwargs)
    131             self._postfit(Y, T, *args, **kwargs)
    132             if inference is not None:

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/econml/iv/nnet/_deepiv.py in fit(self, Y, T, X, Z, inference)
    337         n_components = self._n_components
    338 
--> 339         treatment_network = self._m(z_in, x_in)
    340 
    341         # the dimensionality of the output of the network

<ipython-input-9-a7333673df03> in <lambda>(z, x)
     17 deepIv = DeepIV(
     18     n_components=10,
---> 19     m=lambda z, x: treatment_model(tf.keras.layers.concatenate([z, x])),
     20     h=lambda t, x: hmodel(tf.keras.layers.concatenate([t, x])),
     21     n_samples=2

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/tensorflow/python/keras/layers/merge.py in concatenate(inputs, axis, **kwargs)
    929       A tensor, the concatenation of the inputs alongside axis `axis`.
    930   """
--> 931   return Concatenate(axis=axis, **kwargs)(inputs)
    932 
    933 

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    924     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):
    925       return self._functional_construction_call(inputs, args, kwargs,
--> 926                                                 input_list)
    927 
    928     # Maintains info about the `Layer.call` stack.

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)
   1145           # Node connectivity does not special-case the first argument.
   1146           outputs = self._set_connectivity_metadata((inputs,) + args, kwargs,
-> 1147                                                     outputs)
   1148         self._handle_activity_regularization(inputs, outputs)
   1149         self._set_mask_metadata(inputs, outputs, input_masks, True)

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in _set_connectivity_metadata(self, args, kwargs, outputs)
   2573     # `_outbound_nodes` of the layers that produced the inputs to this
   2574     # layer call.
-> 2575     node_module.Node(self, call_args=args, call_kwargs=kwargs, outputs=outputs)
   2576     return outputs
   2577 

~/.pyenv/versions/miniconda3-4.3.30/envs/my-project/lib/python3.7/site-packages/tensorflow/python/keras/engine/node.py in __init__(self, layer, call_args, call_kwargs, outputs)
    102     self.layer._inbound_nodes.append(self)
    103     for kt in self.keras_inputs:
--> 104       inbound_layer = kt._keras_history.layer
    105       if inbound_layer is not None:  # `None` for `Input` tensors.
    106         inbound_layer._outbound_nodes.append(self)

AttributeError: 'tuple' object has no attribute 'layer'

Addendum

I'm wondering if this might be due to some version incompatibility issues but was unable to find any information regarding this on the documentation. Any help would be greatly appreciated!

Start with the DeepIV usage example and the implementation in econml.iv.nnet, especially DeepIV.fit, then reproduce both failures with the reported Keras 2.3.1, TensorFlow 2.3.0, and EconML 0.10.0 versions. Compare the native Keras and tf.keras paths and document or fix the compatibility issue so the example trains successfully and produces treatment effects.