def masked_index_replace(
    target: torch.Tensor,
    indices: torch.LongTensor,
    mask: torch.BoolTensor,
    replace: torch.Tensor,
) -> torch.Tensor:
    """
    The given `indices` in `target` will be will be replaced with corresponding index
    from the `replace` tensor given a `mask`.


    # Parameters

    target : `torch.Tensor`, required.
        A 3 dimensional tensor of shape (batch_size, sequence_length, embedding_dim).
        This is the tensor to be replaced into.
    indices : `torch.LongTensor`, required
        A 2 dimensional tensor of shape (batch_size, num_indices),
        These are the indices that will be replaced in the original tensor.
    mask : `torch.Tensor`, required.
        A 2 dimensional tensor of shape (batch_size, num_indices), mask.sum() == `nonzero_indices`.
    replace : `torch.Tensor`, required.
        A 3 dimensional tensor of shape (batch_size, num_indices, embedding_dim),
        The tensor to perform scatter from.

    # Returns

    replaced_target : `torch.Tensor`
        A tensor with shape (batch_size, sequence_length, embedding_dim) where 'indices'
        are replaced with the corrosponding vector from `replace`
    """
    target = target.clone()
    mask = mask.bool()
    prev_shape = target.size()
    # Shape: (batch_size * num_indices)
    flattened_indices = flatten_and_batch_shift_indices(indices * mask, target.size(1))
    # Shape: (batch_size * sequence_length, embedding_size)
    flattened_target = target.view(-1, target.size(-1))
    # Shape: (nonzero_indices, 1)
    mask = mask.view(-1)
    flattened_target[flattened_indices[mask]] = replace.view(-1, replace.size(-1))[mask]
    # Shape: (batch_size, sequence_length, embedding_dim)
    replaced_target = flattened_target.reshape(prev_shape)
    return replaced_target
The given `indices` in `target` will be will be replaced with corresponding index
from the `replace` tensor given a `mask`.


# Parameters

target : `torch.Tensor`, required.
A 3 dimensional tensor of shape (batch_size, sequence_length, embedding_dim).
This is the tensor to be replaced into.
indices : `torch.LongTensor`, required
A 2 dimensional tensor of shape (batch_size, num_indices),
These are the indices that will be replaced in the original tensor.
mask : `torch.Tensor`, required.
A 2 dimensional tensor of shape (batch_size, num_indices), mask.sum() == `nonzero_indices`.
replace : `torch.Tensor`, required.
A 3 dimensional tensor of shape (batch_size, num_indices, embedding_dim),
The tensor to perform scatter from.

# Returns

replaced_target : `torch.Tensor`
A tensor with shape (batch_size, sequence_length, embedding_dim) where 'indices'
are replaced with the corrosponding vector from `replace`
allennlp/nn/util.py