title: How to get some layers from trained model and put them into another to make predictions in tensorflow keras
body: <p>I'm having problems using functional api for estimating <img src="https://latex.codecogs.com/gif.latex?X_c%20%5Csim%20%5Cmathcal%7BN%7D%28%20%5CLambda_%7Bab%7D%5E%5Ctop%28x_a%20-%20x_b%29%20+%20%5CLambda_%7Bb%7D%5E%5Ctop%20x_b%2C%20%5Csigma%20%29" alt="normal" /> by maximizing <img src="https://latex.codecogs.com/gif.latex?P%28X_c%7CX_a%2C%20X_b%29%20%3D%20%5Cfrac%7B1%7D%7BZ%7D%5Cexp%7B%5Cleft%5C%7B%5Cfrac%7B-%5C%7C%5CLambda_%7Bb%7D%5E%5Ctop%20x_c%20-%20%5B%5CLambda_%7Bab%7D%5E%5Ctop%28x_a%20-%20x_b%29%20+%20%5CLambda_%7Bb%7D%5E%5Ctop%20x_b%5D%5C%7C%5E2%7D%7B2%5Csigma%5E2%7D%5Cright%5C%7D%7D" alt="proba" /></p>
  <p>First I minimize the error vector by maximizing the probability layer loss, and then I want to use the mean vector layer to rank <code>xc_hat</code> similar embeddings.</p>
  <p>The code is as follows:</p>
  <pre><code>import random as rdn
  import tensorflow.compat.v2 as tf
  tf.enable_v2_behavior()
  import tensorflow_probability as tfp
  import numpy as np
  from sklearn.datasets import make_blobs
  from sklearn.model_selection import train_test_split
  
  tfd = tfp.distributions
  
  n_observations = 2000
  n_features = 5
  d_dim = 3
  lr = 0.005
  
  # Generate toy data
  def make_relations(x_tr, y_tr, c_tr):
    # Generate input data centers being labels (xc) for two random 
    # in-cluster (xa, xb)
    xa = []
    xc = []
    xb = []
  
    for l in y_tr:
      kone = [k for k, lab in zip(x_tr, y_tr) if lab==l]
      if len(kone) &lt; 3:
        continue
      for i, x in enumerate(kone):
        if np.isclose(x, c_tr[l]).all():
            continue
        kone_minus_x = kone.copy()
        kone_minus_x.pop(i)
        print
        xa.append(x) 
        xc.append(c_tr[l])
        xb.append(rdn.choice(kone_minus_x))
  
    return np.vstack(xa), np.vstack(xb), np.vstack(xc)
  
  
  X, Y, C = make_blobs(n_samples=n_observations, 
                                        n_features=n_features, 
                                        centers=int(n_observations*0.2),
                                         return_centers=True)
  x_a, x_b, x_c = make_relations(X, Y, C)
  
  Xa_train, Xa_test = train_test_split(x_a, test_size=.4)
  Xb_train, Xb_test = train_test_split(x_b, test_size=.4)
  Xc_train, Xc_test = train_test_split(x_c, test_size=.4)
  
  Xa_train = Xa_train[np.newaxis]
  Xb_train = Xb_train[np.newaxis]
  Xc_train = Xc_train[np.newaxis]
  
  Xa_test = Xa_test[np.newaxis]
  Xb_test = Xb_test[np.newaxis]
  Xc_test = Xc_test[np.newaxis]
  
  neg_log_likelihood = lambda y, rv_y: -rv_y.log_prob(y)
  
  ones_train = tf.keras.backend.ones((1, Xc_train.shape[1], d_dim)).numpy()
  ones_test = tf.keras.backend.ones((1, Xc_test.shape[1], d_dim)).numpy()
  
  # Build model.
  xa_xb = tf.keras.layers.Input(shape=(None, n_features), name='Xa-Xb')
  L_xa_xb = tf.keras.layers.Dense(d_dim, activation='sigmoid', name='L_Xa-Xb')(xa_xb)
  
  xb = tf.keras.layers.Input(shape=(None, n_features), name='Xb')
  L_xb = tf.keras.layers.Dense(d_dim, activation='sigmoid', name='L_Xb')(xb)
  mu = tf.keras.layers.Add(name='mean_vector')([L_xa_xb, L_xb])
  
  xc = tf.keras.layers.Input(shape=(None, n_features), name='Xc')
  L_xc = tf.keras.layers.Dense(d_dim, name='L_Xc')(xc)
  
  error_vector = tf.keras.layers.Subtract(name='error_vector')([L_xc, mu])
  
  p_xc_given_xa_xb = tfp.layers.DistributionLambda(
      lambda t: tfd.Normal(loc=t, scale=tf.exp(t)), name='Gaussian')(error_vector)
  
  model = tf.keras.Model(inputs=[xa_xb, xb, xc],
                         outputs=p_xc_given_xa_xb, name=&quot;inner_model&quot;)
  model.compile(
        optimizer=tf.optimizers.Adam(learning_rate=lr),
        loss=neg_log_likelihood)
  model.fit([Xa_train - Xb_train, Xb_train, Xc_train], ones_train,
            validation_data=([Xa_test - Xb_test, Xb_test, Xc_test], ones_test),
            epochs=1000,
            verbose=True)
      
  # After trained rebuild the part of the model I will use for prediction
  
  xa_xb = model.get_layer('Xa-Xb')
  L_xa_xb = model.get_layer('L_Xa-Xb')(xa_xb)
  
  xb = model.get_layer('Xb')
  L_xb = model.get_layer('L_Xb')(xb)
  xc = model.get_layer('mean_vector')([L_xa_xb, L_xb])
  
  model = tf.keras.Model(inputs=[xa_xb, xb],
                         outputs=xc, name=&quot;inner_model&quot;)
  
  xc_hat = model([Xa_test - Xb_test, Xb_test])
  </code></pre>
  <p>The idea is to estimate <code>xc</code> However, I have the following error:</p>
  <pre><code>---------------------------------------------------------------------------
  
  AttributeError                            Traceback (most recent call last)
  
  &lt;ipython-input-57-d94b1e8a583c&gt; in &lt;module&gt;()
        2 
        3 xa_xb = model.get_layer('Xa-Xb')
  ----&gt; 4 L_xa_xb = model.get_layer('L_Xa-Xb')(xa_xb)
        5 
        6 xb = model.get_layer('Xb')
  
  1 frames
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
      966           with base_layer_utils.autocast_context_manager(
      967               self._compute_dtype):
  --&gt; 968             outputs = self.call(cast_inputs, *args, **kwargs)
      969           self._handle_activity_regularization(inputs, outputs)
      970           self._set_mask_metadata(inputs, outputs, input_masks)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/core.py in call(self, inputs)
     1178 
     1179   def call(self, inputs):
  -&gt; 1180     rank = inputs.shape.rank
     1181     if rank is not None and rank &gt; 2:
     1182       # Broadcasting is required for the inputs.
  
  AttributeError: 'InputLayer' object has no attribute 'shape'
  </code></pre>
  <p>Im usign Google Colaboratory</p>