def compute_mask(self, inputs, mask=None):
if self.mode == self.MODE_EXPAND:
if self.mask_zero:
output_mask = K.not_equal(inputs, self.mask_zero)
else:
output_mask = None
else:
output_mask = mask
return output_mask
def call(self, inputs, **kwargs):
inputs, tasks = inputs
if K.dtype(tasks) != 'int32':
tasks = K.cast(tasks, 'int32')
task_embed = K.gather(self.embeddings, tasks)
if self.mask_zero:
task_embed = task_embed * K.expand_dims(
K.cast(K.not_equal(tasks, 0), K.floatx()), axis=-1)
return inputs + task_embed
def call(self, inputs, **kwargs):
inputs, tasks = inputs
if K.dtype(tasks) != 'int32':
tasks = K.cast(tasks, 'int32')
task_embed = K.gather(self.embeddings, tasks)
if self.mask_zero:
task_embed = task_embed * K.expand_dims(
K.cast(K.not_equal(tasks, 0), K.floatx()), axis=-1)
if K.backend() == 'theano':
task_embed = K.tile(task_embed, (1, K.shape(inputs)[1], 1))
return inputs + task_embed
def compute_mask(self, inputs, mask=None):
token_mask = K.not_equal(inputs[1], 0)
masked = K.all(K.stack([token_mask, mask[0]], axis=0), axis=0)
if self.return_masked:
return [masked, None]
return masked