'Returns an initializer for the "W_x" parameter matrix.
See equation (23) in the paper. The "W_x" parameter matrix is the
concatenation of two parameter submatrices. The matrix returned is
[W_z, W_r].
Args:
input_dim: The dimension of the cell inputs.
Returns:
A Tensor with shape [input_dim, 2 * num_units] as described above.'

def _w_x_initializer(self, input_dim):

    def _initializer(shape, dtype=tf.float32, partition_info=None):
        num_units = self._num_units
        assert (shape == [input_dim, (2 * num_units)])
        w_z = self._w_initializer([input_dim, num_units], dtype, partition_info)
        w_r = self._w_initializer([input_dim, num_units], dtype, partition_info)
        return tf.concat([w_z, w_r], 1)
    return _initializer
