def __init__(self, units,
kernel_initializer=init.GlorotUniformInitializer(),
bias_initializer=init.RandomNormalInitializer(1e-6)):
super(Dense, self).__init__()
self._units = units
self._kernel_initializer = kernel_initializer
self._bias_initializer = bias_initializer
def __init__(self,
n_heads=1,
d_model=1024,
kernel_initializer=init.GlorotUniformInitializer()):
super(ComputeAttentionOutput, self).__init__()
self._n_heads = n_heads
self._d_model = d_model
self._kernel_initializer = kernel_initializer
def __init__(self,
n_heads=1,
d_head=64,
kernel_initializer=init.GlorotUniformInitializer()):
super(ComputeAttentionHeads, self).__init__()
self._n_heads = n_heads
self._d_head = d_head
self._kernel_initializer = kernel_initializer
def __init__(self,
d_feature,
vocab_size,
kernel_initializer=init.GlorotUniformInitializer()):
super(Embedding, self).__init__()
self._d_feature = d_feature # feature dimensionality
self._vocab_size = vocab_size
self._kernel_initializer = kernel_initializer
def __init__(self,
feature_depth,
vocab_size,
kernel_initializer=init.GlorotUniformInitializer()):
super(Embedding, self).__init__()
self._feature_depth = feature_depth
self._vocab_size = vocab_size
self._kernel_initializer = kernel_initializer
def test_glorot_uniform(self):
initializer = initializers.GlorotUniformInitializer()
input_shape = (29, 5, 7, 20)
init_value = initializer(input_shape, random.get_prng(0))
self.assertEqual(tuple(init_value.shape), input_shape)
