TypeError：一元的坏操作数类型 - ：'ReplicatedVariable'，当我使用tf.keras.layers.PReLU时

我正在转换一个将在GPU上运行的SRGAN keras模型，以便在TPU上运行。

SRGAN使用PReLU激活，当我在我的功能API中使用tf.keras.layers.PReLU时

x = layers.PReLU(alpha_initializer = 'zeros', alpha_regularizer = None, alpha_constraint = None, shared_axes=[1,2])(x)

当我使用TPU调整模型时，我得到以下TypeError

TypeError: bad operand type for unary -: 'ReplicatedVariable'

如果我用以下方式交换PReLU层：

x = tf.Activations('relu')(x)

错误消失了。

有没有人见过这个问题？我认为它可能与PReLU类的调用函数有关：

@tf_export('keras.layers.PReLU')
Class PReLU(layer):
...
  def call(self, inputs, mask=None):
  ...
  else:
  neg = -self.alpha * K.relu(-inputs)