def SumOfWeights():
"""Returns a layer that computes sum of weights."""
return cb.Serial(
cb.Drop(), # Drop inputs.
cb.Drop(), # Drop targets.
core.Sum(axis=None) # Sum weights.
)
def SumOfWeights():
"""Returns a layer to compute sum of weights of all non-masked elements."""
return cb.Serial(
cb.Drop(), # Drop inputs.
cb.Drop(), # Drop targets.
core.Sum(axis=None) # Sum weights.
)
def CountWeights(mask_id=None, has_weights=False):
"""Sum the weights assigned to all elements."""
if has_weights:
return cb.Serial(
cb.Drop(), # Drop inputs.
WeightMask(mask_id=mask_id), # pylint: disable=no-value-for-parameter
cb.Multiply(), # Multiply with provided mask.
core.Sum(axis=None) # Sum all weights.
)
return cb.Serial(
cb.Drop(), # Drop inputs.
WeightMask(mask_id=mask_id), # pylint: disable=no-value-for-parameter
core.Sum(axis=None) # Sum all weights.
)
def SRU(n_units, activation=None):
"""SRU layer as in https://arxiv.org/abs/1709.02755.
As defined in the paper:
(1) y_t = W x_t (+ B optionally, which we do)
(2) f_t = sigmoid(Wf x_t + bf)
(3) r_t = sigmoid(Wr x_t + br)
(4) c_t = f_t * c_{t-1} + (1 - f_t) * y_t
(5) h_t = r_t * activation(c_t) + (1 - r_t) * x_t
We assume the input is of shape [batch, length, depth] and recurrence
happens on the length dimension. This returns a single layer. It's best
to use at least 2, they say in the paper, except inside a Transformer.
Args:
n_units: output depth of the SRU layer.
activation: Optional activation function.
Returns:
The SRU layer.
"""
activation = activation or []
return cb.Serial(
cb.Dup(), # x, x
core.Dense(3 * n_units),
cb.Split(n_items=3), # r, f, y, x
cb.Parallel(core.Sigmoid(), core.Sigmoid()), # r, f, y, x
base.Fn(lambda r, f, y: (y * (1.0 - f), f, r)), # y * (1 - f), f, r, x
cb.Parallel([], [], [cb.Dup(), MakeZeroState()]), # pylint: disable=no-value-for-parameter
cb.Scan(InnerSRUCell(), axis=1), # pylint: disable=no-value-for-parameter
cb.Parallel(activation, cb.Drop()), # act(c), r, x
base.Fn(lambda c, r, x: c * r + x * (1 - r)))
def SumOfWeights(id_to_mask=None, has_weights=False):
"""Returns a layer to compute sum of weights of all non-masked elements."""
multiply_by_weights = cb.Multiply() if has_weights else []
return cb.Serial(
cb.Drop(), # Drop inputs.
_ElementMask(id_to_mask=id_to_mask),
multiply_by_weights,
core.Sum(axis=None) # Sum all.
)
def test_drop(self):
layer = cb.Drop()
input_signature = ShapeDtype((3, 2))
expected_shape = ()
output_shape = base.check_shape_agreement(layer, input_signature)
self.assertEqual(output_shape, expected_shape)
