DeepIV couldn't perform inference?

Hi, I am working on a deep IV estimator, but I couldn't obtain inference on the marginal effect. Below is the testing code adapted from the test_deepiv_models(self) function in the "test_deepiv.py" file.

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

losses = []
marg_effs = []

z_fresh = np.random.uniform(size=(n, 1))
e_fresh = np.random.uniform(low=-0.5, high=0.5, size=(n, 1))
x_fresh = np.random.uniform(size=(n, 1)) + e_fresh
p_fresh = x_fresh + z_fresh * e_fresh + np.random.uniform(size=(n, 1))
y_fresh = p_fresh * x_fresh + e_fresh

for (n1, u, n2) in [(2, False, None), (2, True, None), (1, False, 1)]:
    treatment_model = keras.Sequential([keras.layers.Dense(10, activation='relu', input_shape=(2,)),
                                        keras.layers.Dense(10, activation='relu'),
                                        keras.layers.Dense(10, activation='relu')])

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

    deepIv = DeepIVEstimator(10,
                             lambda z, x: treatment_model(keras.layers.concatenate([z, x])),
                             lambda t, x: hmodel(keras.layers.concatenate([t, x])),
                             n_samples=n1, use_upper_bound_loss=u, n_gradient_samples=n2,
                             first_stage_options={'epochs': epochs}, second_stage_options={'epochs': epochs})
    deepIv.fit(y, p, X=x, Z=z, inference='bootstrap')

The code works well without inference='bootstrap'. But the following error arises when I tried to obtain inference:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-39-a45213b52578> in <module>
     31                              n_samples=n1, use_upper_bound_loss=u, n_gradient_samples=n2,
     32                              first_stage_options={'epochs': epochs}, second_stage_options={'epochs': epochs})
---> 33     deepIv.fit(y, p, X=x, Z=z, inference='bootstrap')
     34 
     35     losses.append(np.mean(np.square(y_fresh - deepIv.predict(p_fresh, x_fresh))))

~/anaconda3/envs/base2/lib/python3.6/site-packages/econml/utilities.py in m(*args, **kwargs)
   1210             if wrong_args:
   1211                 warn(message, category, stacklevel=2)
-> 1212             return to_wrap(*args, **kwargs)
   1213         return m
   1214     return decorator

~/anaconda3/envs/base2/lib/python3.6/site-packages/econml/cate_estimator.py in call(self, Y, T, inference, *args, **kwargs)
    105             if inference is not None:
    106                 # NOTE: we call inference fit *after* calling the main fit method
--> 107                 inference.fit(self, Y, T, *args, **kwargs)
    108             self._inference = inference
    109             return self

~/anaconda3/envs/base2/lib/python3.6/site-packages/econml/inference.py in fit(self, estimator, *args, **kwargs)
     66     def fit(self, estimator, *args, **kwargs):
     67         est = BootstrapEstimator(estimator, self._n_bootstrap_samples, self._n_jobs, compute_means=False,
---> 68                                  bootstrap_type=self._bootstrap_type)
     69         est.fit(*args, **kwargs)
     70         self._est = est

~/anaconda3/envs/base2/lib/python3.6/site-packages/econml/bootstrap.py in __init__(self, wrapped, n_bootstrap_samples, n_jobs, compute_means, bootstrap_type)
     52 
     53     def __init__(self, wrapped, n_bootstrap_samples=1000, n_jobs=None, compute_means=True, bootstrap_type='pivot'):
---> 54         self._instances = [clone(wrapped, safe=False) for _ in range(n_bootstrap_samples)]
     55         self._n_bootstrap_samples = n_bootstrap_samples
     56         self._n_jobs = n_jobs

~/anaconda3/envs/base2/lib/python3.6/site-packages/econml/bootstrap.py in <listcomp>(.0)
     52 
     53     def __init__(self, wrapped, n_bootstrap_samples=1000, n_jobs=None, compute_means=True, bootstrap_type='pivot'):
---> 54         self._instances = [clone(wrapped, safe=False) for _ in range(n_bootstrap_samples)]
     55         self._n_bootstrap_samples = n_bootstrap_samples
     56         self._n_jobs = n_jobs

~/anaconda3/envs/base2/lib/python3.6/site-packages/sklearn/utils/validation.py in inner_f(*args, **kwargs)
     70                           FutureWarning)
     71         kwargs.update({k: arg for k, arg in zip(sig.parameters, args)})
---> 72         return f(**kwargs)
     73     return inner_f
     74 

~/anaconda3/envs/base2/lib/python3.6/site-packages/sklearn/base.py in clone(estimator, safe)
     69     elif not hasattr(estimator, 'get_params') or isinstance(estimator, type):
     70         if not safe:
---> 71             return copy.deepcopy(estimator)
     72         else:
     73             if isinstance(estimator, type):

~/anaconda3/envs/base2/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

~/anaconda3/envs/base2/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    278     if state is not None:
    279         if deep:
--> 280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
    282             y.__setstate__(state)

~/anaconda3/envs/base2/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    148     copier = _deepcopy_dispatch.get(cls)
    149     if copier:
--> 150         y = copier(x, memo)
    151     else:
    152         try:

~/anaconda3/envs/base2/lib/python3.6/copy.py in _deepcopy_dict(x, memo, deepcopy)
    238     memo[id(x)] = y
    239     for key, value in x.items():
--> 240         y[deepcopy(key, memo)] = deepcopy(value, memo)
    241     return y
    242 d[dict] = _deepcopy_dict

~/anaconda3/envs/base2/lib/python3.6/copy.py in deepcopy(x, memo, _nil)
    178                     y = x
    179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
    181 
    182     # If is its own copy, don't memoize.

~/anaconda3/envs/base2/lib/python3.6/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    280             state = deepcopy(state, memo)
    281         if hasattr(y, '__setstate__'):
--> 282             y.__setstate__(state)
    283         else:
    284             if isinstance(state, tuple) and len(state) == 2:

~/.local/lib/python3.6/site-packages/keras/engine/network.py in __setstate__(self, state)
   1332 
   1333     def __setstate__(self, state):
-> 1334         model = saving.unpickle_model(state)
   1335         self.__dict__.update(model.__dict__)
   1336 

~/.local/lib/python3.6/site-packages/keras/engine/saving.py in unpickle_model(state)
    602 def unpickle_model(state):
    603     h5dict = H5Dict(state, mode='r')
--> 604     return _deserialize_model(h5dict)
    605 
    606 

~/.local/lib/python3.6/site-packages/keras/engine/saving.py in _deserialize_model(h5dict, custom_objects, compile)
    272         raise ValueError('No model found in config.')
    273     model_config = json.loads(model_config.decode('utf-8'))
--> 274     model = model_from_config(model_config, custom_objects=custom_objects)
    275     model_weights_group = h5dict['model_weights']
    276 

~/.local/lib/python3.6/site-packages/keras/engine/saving.py in model_from_config(config, custom_objects)
    625                         '`Sequential.from_config(config)`?')
    626     from ..layers import deserialize
--> 627     return deserialize(config, custom_objects=custom_objects)
    628 
    629 

~/.local/lib/python3.6/site-packages/keras/layers/__init__.py in deserialize(config, custom_objects)
    166                                     module_objects=globs,
    167                                     custom_objects=custom_objects,
--> 168                                     printable_module_name='layer')

~/.local/lib/python3.6/site-packages/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    145                     config['config'],
    146                     custom_objects=dict(list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 147                                         list(custom_objects.items())))
    148             with CustomObjectScope(custom_objects):
    149                 return cls.from_config(config['config'])

~/.local/lib/python3.6/site-packages/keras/engine/network.py in from_config(cls, config, custom_objects)
   1073                         node_data = node_data_list[node_index]
   1074                         try:
-> 1075                             process_node(layer, node_data)
   1076 
   1077                         # If the node does not have all inbound layers

~/.local/lib/python3.6/site-packages/keras/engine/network.py in process_node(layer, node_data)
   1023             # and building the layer if needed.
   1024             if input_tensors:
-> 1025                 layer(unpack_singleton(input_tensors), **kwargs)
   1026 
   1027         def process_layer(layer_data):

~/.local/lib/python3.6/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)

~/.local/lib/python3.6/site-packages/keras/engine/base_layer.py in __call__(self, inputs, **kwargs)
    487             # Actually call the layer,
    488             # collecting output(s), mask(s), and shape(s).
--> 489             output = self.call(inputs, **kwargs)
    490             output_mask = self.compute_mask(inputs, previous_mask)
    491 

~/.local/lib/python3.6/site-packages/keras/layers/core.py in call(self, inputs, mask)
    714         else:
    715             self._input_dtypes = K.dtype(inputs)
--> 716         return self.function(inputs, **arguments)
    717 
    718     def compute_mask(self, inputs, mask=None):

~/.local/lib/python3.6/site-packages/keras/layers/core.py in <lambda>(tx)
    375             return K.reshape(all_grads, (-1, d_y, d_t))
    376 
--> 377         self._marginal_effect_model = Model([t_in, x_in], L.Lambda(lambda tx: calc_grad(*tx))([t_in, x_in]))
    378 
    379     def effect(self, X=None, T0=0, T1=1):

TypeError: 'str' object is not callable

Does anyone have similar problems? Is there any solution to it? Thank you in advance!

Start with the adapted test_deepiv.py case and DeepIVEstimator.fit(..., inference='bootstrap'). Trace the failure through inference.py and bootstrap.py, focusing on estimator cloning and the Keras model state shown in the traceback. Done means the supplied DeepIV configuration can fit with bootstrap inference and the marginal-effect inference path no longer raises the reported TypeError.