title: Wrap python callable in keras layer
body: <p>In keras / tensorflow it is often quite simple to describe layers directly as functions that map their input to an output, like so:</p>
  <pre class="lang-py prettyprint-override"><code>def resnet_block(x, kernel_size):
     ch = x.shape[-1]
     out = Conv2D(ch, kernel_size, strides = (1,1), padding='same', activation='relu')(x)
     out = Conv2D(ch, kernel_size, strides = (1,1), padding='same', activation='relu')(out)
     out = Add()([x,out])
     return out
  </code></pre>
  <p>whereas subclassing <code>Layer</code> to get something like</p>
  <pre class="lang-py prettyprint-override"><code>r = ResNetBlock(kernel_size=(3,3))
  y = r(x)
  </code></pre>
  <p>is a little more cumbersome (or even a lot more cumbersome for more complex examples).</p>
  <p>Since keras seems perfectly happy to construct the underlying weights of its layers when they're being called for the first time, I was wondering if it was possible to just wrap functions such as the one above and let keras figure things out once there are inputs, i.e. I would like it to look like this:</p>
  <pre class="lang-py prettyprint-override"><code>r = FunctionWrapperLayer(lambda x:resnet_block(x, kernel_size=(3,3)))
  y = r(x)
  </code></pre>
  <p>I've made an attempt at implementing <code>FunctionWrapperLayer</code>, which looks as follows:</p>
  <pre class="lang-py prettyprint-override"><code>class FunctionWrapperLayer(Layer):
      def __init__(self, fn):
          super(FunctionWrapperLayer, self).__init__()
          self.fn = fn
      
      def build(self, input_shape):
          shape = input_shape[1:]
          inputs = Input(shape)
          outputs = self.fn(inputs)
          self.model = Model(inputs=inputs, outputs=outputs)
          self.model.compile()
          
      def call(self, x):
          return self.model(x)
  </code></pre>
  <p>This looks like it might work, however I've run into some bizarre issues whenever I use activations, e.g. with</p>
  <pre class="lang-py prettyprint-override"><code>def bad(x):
      out = tf.keras.activations.sigmoid(x)
      out = Conv2D(1, (1,1), strides=(1,1), padding='same')(out)
      return out
  x = tf.constant(tf.reshape(tf.range(48,dtype=tf.float32),[1,4,-1,1])
  w = FunctionWrapperLayer(bad)
  w(x)
  </code></pre>
  <p>I get the following error</p>
  <pre><code>FailedPreconditionError:  Error while reading resource variable _AnonymousVar34 from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/_AnonymousVar34/class tensorflow::Var does not exist.
       [[node conv2d_6/BiasAdd/ReadVariableOp (defined at &lt;ipython-input-33-fc380d9255c5&gt;:12) ]] [Op:__inference_keras_scratch_graph_353]
  </code></pre>
  <p>What this suggests to me is that there is something inherently wrong with initializing models like that in the build method. Maybe someone has a better idea as to what might be going on there or how else to get the functionality I would like.</p>
  <p>Update:
  As mentioned by jr15, the above does work when the function involved only uses keras layers. However, the following ALSO works, which has me a little puzzled:</p>
  <pre class="lang-py prettyprint-override"><code>i = Input(x.shape[1:])
  o = bad(i)
  model = Model(inputs=i, outputs=o)
  model(x)
  </code></pre>
  <p>Incidentally, <code>model.submodules</code> yields</p>
  <pre><code>(&lt;tensorflow.python.keras.engine.input_layer.InputLayer at 0x219d80c77c0&gt;,
   &lt;tensorflow.python.keras.engine.base_layer.TensorFlowOpLayer at 0x219d7afc820&gt;,
   &lt;tensorflow.python.keras.layers.convolutional.Conv2D at 0x219d7deafa0&gt;)
  </code></pre>
  <p>meaning the activation is automatically turned into a &quot;TensorFlowOpLayer&quot; when doing it like that.</p>
  <p>Another update:
  Looking at the original error message, it seems like the activation isn't the only culprit. If I remove the convolution and use the wrapper everything works as well and again I find a &quot;TensorFlowOpLayer&quot; when inspecting the submodules.</p>