Function masked_scatter

tensorrt_llm/functional.py:2539–2574

Add the masked_scatter base on PyTorch definition. See https://pytorch.org/docs/stable/generated/torch.Tensor.masked_scatter_.html#torch-tensor-masked-scatter for a description of that function. Parameters: input : Tensor The input tensor. mask : Tenso

(input: Tensor, mask: Tensor, source: Tensor)

def masked_scatter(input: Tensor, mask: Tensor, source: Tensor) -> Tensor:
    '''
    Add the masked_scatter base on PyTorch definition.

    See https://pytorch.org/docs/stable/generated/torch.Tensor.masked_scatter_.html#torch-tensor-masked-scatter for a
    description of that function.

    Parameters:
        input : Tensor
            The input tensor.

        mask : Tensor
            The boolean mask tensor that indicates elements to select.

        source: Tensor
            The tensor to copy from
    Returns:
        The tensor containing the source tensor selected by mask.

    '''
    assert input.rank() >= 1, "input should have rank >= 1"
    input, mask = broadcast_helper(input, mask)
    expanded_mask = expand(mask, shape(input))

    non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor)

    shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0))
    shuffle_layer.second_transpose = (1, 0)
    source = source.view([-1])

    scatter_layer = default_trtnet().add_scatter(input.trt_tensor,
                                                 shuffle_layer.get_output(0),
                                                 source.trt_tensor,
                                                 mode=trt.ScatterMode.ND)

    return _create_tensor(scatter_layer.get_output(0), scatter_layer)


def concat(inputs: Sequence[Union[Tensor, int]], dim: int = 0) -> Tensor: