def _step_fn(inputs):
cache = _create_cache(batch_size, from_seq_length,
num_heads, head_size)
mask = t5.make_causal_mask(tf.ones((batch_size, 1)))
return l(query=inputs,
mask=mask,
cache=cache,
decode_position=decode_position)
def _step_fn(inputs):
cache = _create_cache(batch_size, from_seq_length,
num_heads, head_size)
mask = t5.make_causal_mask(tf.ones((batch_size, 1)))
position_bias = pos_embed(from_seq_length, from_seq_length)
return l(hidden_states=inputs,
cache=cache,
attention_mask=mask,
decode_position=decode_position,
position_bias=position_bias)
def test_masks(self):
causal_mask = t5.make_causal_mask(np.zeros((2, 5)))
self.assertEqual(causal_mask.shape, (2, 1, 5, 5))