<p dir="auto">I want to implement a transformer toolkit(a set of modules that will compose some kind of DL model with attention ).I have said a toolkit because the today's landscape of the transformer research is very diverse and I think that it will be better to have the separate modules ( attention layer , concat layer ,add&amp; norm layer etc)  and mix and match them rather to have one or two ready layer that will be very general and not good for any specific task</p>
<p dir="auto">I'm trying to use a tf op inside a Lambda layer. The goal is to basically wrap tf layers such that <code class="notranslate">model.compile</code> etc. still work. As I understood it I can just incorporate tf ops inside Lambda layers. If I do this however, serialization seems to be broken because it fails to pickle:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="TypeError: can't pickle _thread.lock objects"><pre class="notranslate"><code class="notranslate">TypeError: can't pickle _thread.lock objects
</code></pre></div>
<p dir="auto">The weird thing is that if I create a model in the main method, it seems to work fine. If I create the model inside a separate method, it fails to serialize. I'm guessing this is due to some scoping problem, but it could also be a Keras issue. A minimum example where this occurs is shown below, where I replaced a <code class="notranslate">Conv2D</code> with a tf <code class="notranslate">conv2d</code> op to demonstrate the issue.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import tensorflow as tf

from keras.models import Model
from keras.layers import Input, Lambda

def weight_variable(shape):
	initial = tf.truncated_normal(shape, stddev=0.1)
	return tf.Variable(initial)

def conv2d(x, W):
	return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def create_model(num_classes):
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)
	model = Model(inputs=inputs, outputs=x)
	return model

if __name__=='__main__':
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)

	# This works
	model_working = Model(inputs=inputs, outputs=x)
	print(model_working.to_json())

	# This gives &quot;TypeError: can't pickle _thread.lock objects&quot;
	model_broken = create_model(10)
	print(model_broken.to_json())"><pre class="notranslate"><code class="notranslate">import tensorflow as tf

from keras.models import Model
from keras.layers import Input, Lambda

def weight_variable(shape):
	initial = tf.truncated_normal(shape, stddev=0.1)
	return tf.Variable(initial)

def conv2d(x, W):
	return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def create_model(num_classes):
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)
	model = Model(inputs=inputs, outputs=x)
	return model

if __name__=='__main__':
	inputs = Input(shape=(28, 28, 1))
	W_conv1 = weight_variable([3, 3, 1, 32])
	x = Lambda(lambda x: conv2d(x, W_conv1))(inputs)

	# This works
	model_working = Model(inputs=inputs, outputs=x)
	print(model_working.to_json())

	# This gives "TypeError: can't pickle _thread.lock objects"
	model_broken = create_model(10)
	print(model_broken.to_json())
</code></pre></div>
<p dir="auto">edit: <code class="notranslate">to_json</code>, <code class="notranslate">to_yaml</code> and <code class="notranslate">save</code> are at least broken with this code. I'm guessing its all methods that require the serialization of the network architecture that are broken with this model.</p>