Apply weight drop to the parameter of a block.

    Parameters
    ----------
    block : Block or HybridBlock
        The block whose parameter is to be applied weight-drop.
    local_param_regex : str
        The regex for parameter names used in the self.params.get(), such as 'weight'.
    rate : float
        Fraction of the input units to drop. Must be a number between 0 and 1.
    axes : tuple of int, default ()
        The axes on which dropout mask is shared. If empty, regular dropout is applied.
    weight_drop_mode : {'training', 'always'}, default 'training'
        Whether the weight dropout should be applied only at training time, or always be applied.

    Examples
    --------
    >>> net = gluon.rnn.LSTM(10, num_layers=2, bidirectional=True)
    >>> gluonnlp.model.apply_weight_drop(net, r'.*h2h_weight', 0.5)
    >>> net.collect_params()
    lstm0_ (
      Parameter lstm0_l0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_l0_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_l0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_r0_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_r0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_l1_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_l1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_r1_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_r1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
    )
    >>> ones = mx.nd.ones((3, 4, 5))
    >>> net.initialize()
    >>> with mx.autograd.train_mode():
    ...     net(ones).max().asscalar() != net(ones).max().asscalar()
    True

def apply_weight_drop(block, local_param_regex, rate, axes=(),
                      weight_dropout_mode='training'):
    """Apply weight drop to the parameter of a block.

    Parameters
    ----------
    block : Block or HybridBlock
        The block whose parameter is to be applied weight-drop.
    local_param_regex : str
        The regex for parameter names used in the self.params.get(), such as 'weight'.
    rate : float
        Fraction of the input units to drop. Must be a number between 0 and 1.
    axes : tuple of int, default ()
        The axes on which dropout mask is shared. If empty, regular dropout is applied.
    weight_drop_mode : {'training', 'always'}, default 'training'
        Whether the weight dropout should be applied only at training time, or always be applied.

    Examples
    --------
    >>> net = gluon.rnn.LSTM(10, num_layers=2, bidirectional=True)
    >>> gluonnlp.model.apply_weight_drop(net, r'.*h2h_weight', 0.5)
    >>> net.collect_params()
    lstm0_ (
      Parameter lstm0_l0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_l0_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_l0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r0_i2h_weight (shape=(40, 0), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_r0_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_r0_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r0_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_l1_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_l1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_l1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r1_i2h_weight (shape=(40, 20), dtype=<class 'numpy.float32'>)
      WeightDropParameter lstm0_r1_h2h_weight (shape=(40, 10), dtype=<class 'numpy.float32'>, \
rate=0.5, mode=training)
      Parameter lstm0_r1_i2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
      Parameter lstm0_r1_h2h_bias (shape=(40,), dtype=<class 'numpy.float32'>)
    )
    >>> ones = mx.nd.ones((3, 4, 5))
    >>> net.initialize()
    >>> with mx.autograd.train_mode():
    ...     net(ones).max().asscalar() != net(ones).max().asscalar()
    True
    """
    if not rate:
        return

    existing_params = _find_params(block, local_param_regex)
    for (local_param_name, param), \
            (ref_params_list, ref_reg_params_list) in existing_params.items():
        dropped_param = WeightDropParameter(param, rate, weight_dropout_mode, axes)
        for ref_params in ref_params_list:
            ref_params[param.name] = dropped_param
        for ref_reg_params in ref_reg_params_list:
            ref_reg_params[local_param_name] = dropped_param
            if hasattr(block, local_param_name):
                local_attr = getattr(block, local_param_name)
                if local_attr == param:
                    local_attr = dropped_param
                elif isinstance(local_attr, (list, tuple)):
                    if isinstance(local_attr, tuple):
                        local_attr = list(local_attr)
                    for i, v in enumerate(local_attr):
                        if v == param:
                            local_attr[i] = dropped_param
                elif isinstance(local_attr, dict):
                    for k, v in local_attr:
                        if v == param:
                            local_attr[k] = dropped_param
                else:
                    continue
                if local_attr:
                    super(Block, block).__setattr__(local_param_name, local_attr)