使用Keras模型的功能构建TPU策略

Question

在模型生成功能下定义模型时，是否可以实例化TPU keras模型？

我希望拥有一个可以在TPU中使用的keras模型，该模型可以放在sk-learn包装器中，并可以在管道中使用，如下所示。

添加： 据我在tensorflow指南上了解到，Keras模型定义需要包含在TPUstrategy.scope（）中，但是我不确定这是否可以泛化为模型生成函数吗？

with tpu_strategy.scope():
  embedding_layer = Embedding(...)

  def model_generator():
    input = Input(...)
    x = embedding_layer(tokens_input)
    x = ... # more layers
    outp = Dense(1, activation="sigmoid")(x)

    model = Model(inputs=[input], outputs=[outp])
    model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
    return model

  model = KerasRegressor(model_generator, epochs=10, batch_size=batch_size, verbose=1)

## elsewhere...
model.fit(X,y)