torch.nested

Введение

Вложенные тензоры позволяют хранить данные с неоднородной формой и выполнять над ними операции как над единым тензором. Такие данные хранятся в эффективном упакованном представлении, при этом для применения операций предоставляется стандартный интерфейс тензоров PyTorch.

Вложенные тензоры часто применяются для представления пакетов последовательных данных переменной длины из разных областей, например предложений разной длины, изображений разных размеров, а также аудио- и видеоклипов разной длительности. Традиционно такие данные обрабатывались путём дополнения последовательностей до максимальной длины в пакете, выполнения вычислений над дополненной формой и последующего применения маски для удаления дополнения. Это неэффективно и чревато ошибками; вложенные тензоры призваны решить эти проблемы.

API для вызова операций над вложенным тензором ничем не отличается от API для обычного torch.Tensor, что обеспечивает бесшовную интеграцию с существующими моделями. Главное отличие заключается в создании входных данных.

Поскольку эта функция является прототипом, набор поддерживаемых операций ограничен, но постоянно расширяется. Мы будем рады сообщениям о проблемах, запросам на добавление функций и вкладу в разработку. Подробнее об участии в разработке можно узнать в этом файле Readme.

Создание

Создать вложенный тензор просто: достаточно передать список тензоров конструктору torch.nested.nested_tensor. Вложенный тензор с компоновкой torch.jagged (также называемый «NJT») поддерживает одно измерение с неоднородной длиной. Этот конструктор копирует входные тензоры в упакованный непрерывный блок памяти в соответствии с компоновкой, описанной в разделе data_layout_ ниже.

>>> a, b = torch.arange(3), torch.arange(5) + 3
>>> a
tensor([0, 1, 2])
>>> b
tensor([3, 4, 5, 6, 7])
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> print([component for component in nt])
[tensor([0, 1, 2]), tensor([3, 4, 5, 6, 7])]

Все тензоры в списке должны иметь одинаковое число измерений, но в остальном их формы могут различаться вдоль одного измерения. Если размерности входных компонентов не совпадают, конструктор выдаёт ошибку.

>>> a = torch.randn(50, 128) # 2D tensor
>>> b = torch.randn(2, 50, 128) # 3D tensor
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
...
RuntimeError: When constructing a nested tensor, all tensors in list must have the same dim

При создании можно задать dtype, устройство и необходимость вычисления градиентов с помощью обычных именованных аргументов.

>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32, device="cuda", requires_grad=True)
>>> print([component for component in nt])
[tensor([0., 1., 2.], device='cuda:0',
       grad_fn=<UnbindBackwardAutogradNestedTensor0>), tensor([3., 4., 5., 6., 7.], device='cuda:0',
       grad_fn=<UnbindBackwardAutogradNestedTensor0>)]

Для сохранения истории autograd тензоров, переданных конструктору, можно использовать torch.nested.as_nested_tensor. При использовании этого конструктора градиенты будут проходить через вложенный тензор обратно к исходным компонентам. Обратите внимание, что этот конструктор по-прежнему копирует входные компоненты в упакованный непрерывный блок памяти.

>>> a = torch.randn(12, 512, requires_grad=True)
>>> b = torch.randn(23, 512, requires_grad=True)
>>> nt = torch.nested.as_nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt.sum().backward()
>>> a.grad
tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        ...,
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.]])
>>> b.grad
tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        ...,
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.]])

Все приведённые выше функции создают непрерывные NJT, для хранения упакованной формы базовых компонентов которых выделяется блок памяти (подробнее см. раздел data_layout_ ниже).

Также можно создать неконтинуальное представление NJT поверх уже существующего плотного тензора с дополнением, избежав выделения памяти и копирования. Для этого предназначен torch.nested.narrow().

>>> padded = torch.randn(3, 5, 4)
>>> seq_lens = torch.tensor([3, 2, 5], dtype=torch.int64)
>>> nt = torch.nested.narrow(padded, dim=1, start=0, length=seq_lens, layout=torch.jagged)
>>> nt.shape
torch.Size([3, j1, 4])
>>> nt.is_contiguous()
False

Обратите внимание: вложенный тензор представляет собой представление исходного дополненного плотного тензора и ссылается на ту же память без копирования или выделения памяти. Поддержка операций для неконтинуальных NJT несколько ограничена, поэтому при нехватке поддержки всегда можно преобразовать тензор в непрерывный NJT с помощью contiguous().

Компоновка и форма данных

Для повышения эффективности компоненты вложенных тензоров обычно упаковываются в непрерывный блок памяти, а дополнительные метаданные задают границы элементов пакета. Для компоновки torch.jagged непрерывный блок памяти хранится в компоненте values, а компонент offsets задаёт границы элементов пакета для измерения с неоднородной длиной.

При необходимости можно напрямую обращаться к базовым компонентам NJT.

>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(32, 128) # text 2
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt.values().shape  # note the "packing" of the ragged dimension; no padding needed
torch.Size([82, 128])
>>> nt.offsets()
tensor([ 0, 50, 82])

Также может быть полезно создавать NJT напрямую из составляющих values и offsets для jagged-представления; для этого предназначен конструктор torch.nested.nested_tensor_from_jagged().

>>> values = torch.randn(82, 128)
>>> offsets = torch.tensor([0, 50, 82], dtype=torch.int64)
>>> nt = torch.nested.nested_tensor_from_jagged(values=values, offsets=offsets)

Форма NJT определена однозначно, а его размерность на единицу больше размерности составляющих его компонентов. Структура измерения с неоднородной длиной представлена символическим значением (в приведённом ниже примере — j1).

>>> a = torch.randn(50, 128)
>>> b = torch.randn(32, 128)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt.dim()
3
>>> nt.shape
torch.Size([2, j1, 128])

Чтобы NJT были совместимы друг с другом, их структуры с неоднородной длиной должны совпадать. Например, для выполнения бинарной операции над двумя NJT их структуры с неоднородной длиной должны совпадать (то есть в их формах должен быть один и тот же символ неоднородной формы). На практике каждый символ соответствует конкретному тензору offsets, поэтому для совместимости оба NJT должны иметь один и тот же тензор offsets.

>>> a = torch.randn(50, 128)
>>> b = torch.randn(32, 128)
>>> nt1 = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt2 = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt1.offsets() is nt2.offsets()
False
>>> nt3 = nt1 + nt2
RuntimeError: cannot call binary pointwise function add.Tensor with inputs of shapes (2, j2, 128) and (2, j3, 128)

В приведённом выше примере концептуальные формы двух NJT одинаковы, но они не ссылаются на один и тот же тензор offsets, поэтому их формы различаются и они несовместимы. Мы понимаем, что такое поведение неинтуитивно, и активно работаем над смягчением этого ограничения к бета-версии вложенных тензоров. Способ обхода этого ограничения описан в разделе Устранение неполадок этого документа.

Помимо метаданных offsets, для компонентов NJT также можно вычислять и кэшировать минимальную и максимальную длину последовательностей. Это может пригодиться для вызова определённых ядер (например, SDPA). Сейчас общедоступных API для доступа к этим данным нет, но ситуация изменится к выходу бета-версии.

Поддерживаемые операции

В этом разделе перечислены распространённые операции над вложенными тензорами, которые могут быть полезны. Список не является исчерпывающим: в PyTorch доступны порядка нескольких тысяч операций. Сейчас для вложенных тензоров поддерживается значительная их часть, однако реализация полной поддержки — масштабная задача. В идеале для вложенных тензоров должны поддерживаться все операции PyTorch, доступные для обычных тензоров. Чтобы помочь нам достичь этой цели, рекомендуем:

• Оставить здесь запрос на нужные вам операции, чтобы помочь нам определить приоритеты.
• Принять участие в разработке! Добавить поддержку вложенных тензоров для конкретной операции PyTorch не так сложно; подробности см. в разделе Участие в разработке ниже.

>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(3, 3)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> nt.unbind()
(tensor([[-0.9916, -0.3363, -0.2799],
        [-2.3520, -0.5896, -0.4374]]), tensor([[-2.0969, -1.0104,  1.4841],
        [ 2.0952,  0.2973,  0.2516],
        [ 0.9035,  1.3623,  0.2026]]))
>>> nt.unbind()[0] is not a
True
>>> nt.unbind()[0].mul_(3)
tensor([[ 3.6858, -3.7030, -4.4525],
        [-2.3481,  2.0236,  0.1975]])
>>> nt.unbind()
(tensor([[-2.9747, -1.0089, -0.8396],
        [-7.0561, -1.7688, -1.3122]]), tensor([[-2.0969, -1.0104,  1.4841],
        [ 2.0952,  0.2973,  0.2516],
        [ 0.9035,  1.3623,  0.2026]]))

>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(6, 3)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> padded = torch.nested.to_padded_tensor(nt, padding=4.2)
>>> padded
tensor([[[ 1.6107,  0.5723,  0.3913],
         [ 0.0700, -0.4954,  1.8663],
         [ 4.2000,  4.2000,  4.2000],
         [ 4.2000,  4.2000,  4.2000],
         [ 4.2000,  4.2000,  4.2000],
         [ 4.2000,  4.2000,  4.2000]],
        [[-0.0479, -0.7610, -0.3484],
         [ 1.1345,  1.0556,  0.3634],
         [-1.7122, -0.5921,  0.0540],
         [-0.5506,  0.7608,  2.0606],
         [ 1.5658, -1.1934,  0.3041],
         [ 0.1483, -1.1284,  0.6957]]])

>>> padded = torch.randn(3, 5, 4)
>>> seq_lens = torch.tensor([3, 2, 5], dtype=torch.int64)
>>> nt = torch.nested.narrow(padded, dim=1, length=seq_lens, layout=torch.jagged)
>>> nt.shape
torch.Size([3, j1, 4])
>>> nt = nt.contiguous()
>>> nt.shape
torch.Size([3, j2, 4])

>>> a = torch.randn(2, 6)
>>> b = torch.randn(4, 6)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> nt.shape
torch.Size([2, j1, 6])
>>> nt.unsqueeze(-1).shape
torch.Size([2, j1, 6, 1])
>>> nt.unflatten(-1, [2, 3]).shape
torch.Size([2, j1, 2, 3])
>>> torch.cat([nt, nt], dim=2).shape
torch.Size([2, j1, 12])
>>> torch.stack([nt, nt], dim=2).shape
torch.Size([2, j1, 2, 6])
>>> nt.transpose(-1, -2).shape
torch.Size([2, 6, j1])

Использование с torch.compile

NJT предназначены для использования с torch.compile() для достижения оптимальной производительности, поэтому мы всегда рекомендуем по возможности использовать torch.compile() с NJT. NJT поддерживаются сразу и без разрыва графа как при передаче в качестве входных данных скомпилированной функции или модуля, так и при создании непосредственно внутри функции.

If you're not able to utilize ``torch.compile()`` for your use case, performance and memory
usage may still benefit from the use of NJTs, but it's not as clear-cut whether this will be
the case. It is important that the tensors being operated on are large enough so the
performance gains are not outweighed by the overhead of python tensor subclasses.

>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(4, 3)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> def f(x): return x.sin() + 1
...
>>> compiled_f = torch.compile(f, fullgraph=True)
>>> output = compiled_f(nt)
>>> output.shape
torch.Size([2, j1, 3])
>>> def g(values, offsets): return torch.nested.nested_tensor_from_jagged(values, offsets) * 2.
...
>>> compiled_g = torch.compile(g, fullgraph=True)
>>> output2 = compiled_g(nt.values(), nt.offsets())
>>> output2.shape
torch.Size([2, j1, 3])

Обратите внимание, что NJT поддерживают динамические формы, что позволяет избежать ненужной повторной компиляции при изменении структуры с неоднородной длиной.

>>> a = torch.randn(2, 3)
>>> b = torch.randn(4, 3)
>>> c = torch.randn(5, 3)
>>> d = torch.randn(6, 3)
>>> nt1 = torch.nested.nested_tensor([a, b], layout=torch.jagged)
>>> nt2 = torch.nested.nested_tensor([c, d], layout=torch.jagged)
>>> def f(x): return x.sin() + 1
...
>>> compiled_f = torch.compile(f, fullgraph=True)
>>> output1 = compiled_f(nt1)
>>> output2 = compiled_f(nt2)  # NB: No recompile needed even though ragged structure differs

Если при использовании NJT вместе с torch.compile возникают проблемы или непонятные ошибки, сообщите о них в PyTorch. Полная поддержка подклассов в torch.compile — долгосрочная задача, поэтому сейчас возможны некоторые недоработки.

Устранение неполадок

В этом разделе описаны распространённые ошибки, с которыми можно столкнуться при использовании вложенных тензоров, объясняются причины их возникновения и предлагаются способы их устранения.

    NotImplementedError: aten.view_as_real.default

    RuntimeError: cannot call binary pointwise function add.Tensor with inputs of shapes (2, j2, 128) and (2, j3, 128)

>>> a = torch.randn(50, 128)
>>> b = torch.randn(32, 128)
>>> nt1 = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> nt2 = torch.nested.nested_tensor_from_jagged(values=torch.randn(82, 128), offsets=nt1.offsets())
>>> nt3 = nt1 + nt2
>>> nt3.shape
torch.Size([2, j1, 128])

    torch._dynamo.exc.Unsupported: data dependent operator: aten._local_scalar_dense.default; to enable, set torch._dynamo.config.capture_scalar_outputs = True

>>> a = torch.randn(50, 128)
>>> b = torch.randn(32, 128)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, dtype=torch.float32)
>>> def f(nt): return nt.chunk(2, dim=0)[0]
...
>>> compiled_f = torch.compile(f, fullgraph=True)
>>> output = compiled_f(nt)

>>> torch._dynamo.config.capture_dynamic_output_shape_ops = True
>>> torch._dynamo.config.capture_scalar_outputs = True

Участие в разработке

Если вы хотите принять участие в разработке вложенных тензоров, один из наиболее эффективных способов помочь — добавить поддержку вложенных тензоров для операции PyTorch, которая пока не поддерживается. Обычно для этого нужно выполнить несколько простых шагов:

1. Определите имя операции, поддержку которой нужно добавить; оно должно выглядеть примерно так: aten.view_as_real.default. Сигнатуру этой операции можно найти в aten/src/ATen/native/native_functions.yaml.
2. Зарегистрируйте реализацию операции в torch/nested/_internal/ops.py, следуя приведённому там шаблону для других операций. Для проверки схемы используйте сигнатуру из native_functions.yaml.

Чаще всего операцию реализуют, разворачивая NJT в составляющие его компоненты, повторно отправляя операцию на базовый буфер values и перенося соответствующие метаданные NJT (включая offsets) в новый выходной NJT. Если ожидается, что результат операции будет иметь форму, отличную от формы входных данных, необходимо вычислить новые метаданные offsets и т. д.

При выполнении операции над измерением пакета или измерением с неоднородной длиной эти приёмы помогут быстро создать работающую реализацию:

• Для операции, выполняемой отдельно над каждым элементом пакета, должен подойти резервный вариант на основе unbind().
• Для операции над измерением с неоднородной длиной рассмотрите возможность преобразования данных в дополненный плотный тензор с правильно выбранным значением дополнения, которое не будет искажать результат, выполнения операции и обратного преобразования в NJT. В рамках torch.compile эти преобразования можно объединить, чтобы избежать материализации промежуточного дополненного тензора.

Подробная документация по функциям создания и преобразования

torch.nested.nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False) [исходный код]

Создаёт вложенный тензор без истории Autograd (также известный как «тензор-лист», см. Механика Autograd) из tensor_list списка тензоров.

Параметры:

• tensor_list (List[array_like]) – список тензоров или любых объектов, которые можно передать в torch.tensor,
• dimensionality. (где каждый элемент списка имеет одинаковую размерность) –

Именованные аргументы:

• dtype (torch.dtype, необязательно) – требуемый тип возвращаемого вложенного тензора. По умолчанию: если None, тот же torch.dtype, что и у самого левого тензора в списке.
• layout (torch.layout, необязательно) – требуемая компоновка возвращаемого вложенного тензора. Поддерживаются только страйдовая и рваная компоновки. По умолчанию: если None, страйдовая компоновка.
• device (torch.device, необязательно) – требуемое устройство для возвращаемого вложенного тензора. По умолчанию: если None, то же torch.device, что и у самого левого тензора в списке
• requires_grad (bool, необязательно) – следует ли Autograd записывать операции над возвращаемым вложенным тензором. По умолчанию: False.
• pin_memory (bool, необязательно) – если задано, возвращаемый вложенный тензор будет размещён в закреплённой памяти. Работает только для тензоров CPU. По умолчанию: False.

Тип возвращаемого значения:

Tensor

Пример:

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged, requires_grad=True)
>>> nt.is_leaf
True

torch.nested.nested_tensor_from_jagged(values, offsets=None, lengths=None, jagged_dim=None, min_seqlen=None, max_seqlen=None) [исходный код]

Создаёт вложенный тензор с рваной компоновкой из заданных рваных компонентов. Рваная компоновка состоит из обязательного буфера значений, в котором данные рваного измерения упакованы в одно измерение. Метаданные offsets / lengths определяют, как это измерение разбивается на элементы пакета, и должны быть размещены на том же устройстве, что и буфер значений.

Ожидаемые форматы метаданных:

• offsets: индексы в упакованном измерении, разделяющие его на элементы пакета разного размера. Пример: [0, 2, 3, 6] означает, что упакованное рваное измерение размера 6 концептуально разбивается на элементы пакета длины [2, 1, 3]. Обратите внимание, что для удобства работы ядер требуются смещения как в начале, так и в конце (то есть размер batch_size + 1).
• lengths: длины отдельных элементов пакета; shape == batch_size. Пример: [2, 1, 3] означает, что упакованное рваное измерение размера 6 концептуально разбивается на элементы пакета длины [2, 1, 3].

Обратите внимание, что может быть полезно задать и offsets, и lengths. Это описывает вложенный тензор с «пропусками»: offsets указывают начальную позицию каждого элемента пакета, а length задаёт общее количество элементов (см. пример ниже).

Возвращаемый вложенный тензор с рваной компоновкой будет представлением входного тензора values.

Параметры:

• values (torch.Tensor) – базовый буфер формы (sum_B(*), D_1, …, D_N). Данные рваного измерения упакованы в одно измерение; для разделения элементов пакета используются метаданные offsets / lengths.
• offsets (необязательный torch.Tensor) – смещения в рваном измерении формы B + 1.
• lengths (необязательный torch.Tensor) – длины элементов пакета формы B.
• jagged_dim (необязательный python:int) – указывает, какое измерение в values является упакованным рваным измерением. Должно быть >= 1, поскольку измерение пакета (dim=0) не может быть рваным. Если None, устанавливается dim=1 (то есть измерение непосредственно после измерения пакета). По умолчанию: None
• min_seqlen (необязательный python:int) – если задано, использует указанное значение в качестве кэшированной минимальной длины последовательности для возвращаемого вложенного тензора. Это может быть полезной альтернативой вычислению значения по запросу и, возможно, позволит избежать синхронизации GPU -> CPU. По умолчанию: None
• max_seqlen (необязательный python:int) – если задано, использует указанное значение в качестве кэшированной максимальной длины последовательности для возвращаемого вложенного тензора. Это может быть полезной альтернативой вычислению значения по запросу и, возможно, позволит избежать синхронизации GPU -> CPU. По умолчанию: None

Тип возвращаемого значения:

Tensor

Пример:

>>> values = torch.randn(12, 5)
>>> offsets = torch.tensor([0, 3, 5, 6, 10, 12])
>>> nt = nested_tensor_from_jagged(values, offsets)
>>> # 3D shape with the middle dimension jagged
>>> nt.shape  # xdoctest: +ELLIPSIS
torch.Size([5, j..., 5])
>>> # Length of each item in the batch:
>>> offsets.diff()
tensor([3, 2, 1, 4, 2])

>>> values = torch.randn(6, 5)
>>> offsets = torch.tensor([0, 2, 3, 6])
>>> lengths = torch.tensor([1, 1, 2])
>>> # NT with holes
>>> nt = nested_tensor_from_jagged(values, offsets, lengths)
>>> a, b, c = nt.unbind()
>>> # Batch item 1 consists of indices [0, 1)
>>> torch.equal(a, values[0:1, :])
True
>>> # Batch item 2 consists of indices [2, 3)
>>> torch.equal(b, values[2:3, :])
True
>>> # Batch item 3 consists of indices [3, 5)
>>> torch.equal(c, values[3:5, :])
True

torch.nested.as_nested_tensor(ts, dtype=None, device=None, layout=None) [исходный код]

Создаёт вложенный тензор, сохраняя историю Autograd, из тензора или списка / кортежа тензоров.

Если передан вложенный тензор, он будет возвращён напрямую, если только устройство / тип данных / компоновка не отличаются. Обратите внимание, что преобразование устройства / типа данных приводит к копированию, а преобразование компоновки в этой функции пока не поддерживается.

Если передан невложенный тензор, он рассматривается как пакет элементов одинакового размера. Копирование выполняется, если переданные устройство / тип данных отличаются от устройства / типа данных входных данных ИЛИ если входные данные не являются непрерывными. В противном случае напрямую используется хранилище входных данных.

Если предоставлен список тензоров, при создании вложенного тензора тензоры из списка всегда копируются.

Параметры:

ts (Tensor или List[Tensor] или Tuple[Tensor]) – тензор, который следует рассматривать как вложенный тензор, ИЛИ список / кортеж тензоров с одинаковым ndim

Именованные аргументы:

• dtype (torch.dtype, необязательно) – требуемый тип возвращаемого вложенного тензора. По умолчанию: если None, тот же torch.dtype, что и у самого левого тензора в списке.
• device (torch.device, необязательно) – требуемое устройство для возвращаемого вложенного тензора. По умолчанию: если None, то же torch.device, что и у самого левого тензора в списке
• layout (torch.layout, необязательно) – требуемая компоновка возвращаемого вложенного тензора. Поддерживаются только страйдовая и рваная компоновки. По умолчанию: если None, страйдовая компоновка.

Тип возвращаемого значения:

Tensor

Пример:

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.as_nested_tensor([a, b], layout=torch.jagged)
>>> nt.is_leaf
False
>>> buffer = torch.cat([torch.ones_like(a), torch.zeros_like(b)])
>>> fake_grad = torch.nested.nested_tensor_from_jagged(buffer, nt.offsets())
>>> nt.backward(fake_grad)
>>> a.grad
tensor([1., 1., 1.])
>>> b.grad
tensor([0., 0., 0., 0., 0.])
>>> c = torch.randn(3, 5, requires_grad=True)
>>> nt2 = torch.nested.as_nested_tensor(c, layout=torch.jagged)

torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) → Tensor [исходный код]

Возвращает новый (невложенный) тензор, дополняя input вложенный тензор. Начальные элементы заполняются вложенными данными, а оставшиеся — значениями заполнения.

Предупреждение

to_padded_tensor() всегда копирует базовые данные, поскольку вложенные и невложенные тензоры имеют разную компоновку в памяти.

Параметры:

padding (float) – значение заполнения для оставшихся элементов.

Именованные аргументы:

• output_size (Tuple[int]) – размер выходного тензора. Если задан, он должен быть достаточно большим, чтобы вместить все вложенные данные; в противном случае размер будет определён как максимальный размер каждого вложенного подтензора вдоль каждого измерения.
• out (Tensor, необязательно) – выходной тензор.

Пример:

>>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))])
nested_tensor([
  tensor([[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
          [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995]]),
  tensor([[-1.8546, -0.7194, -0.2918, -0.1846],
          [ 0.2773,  0.8793, -0.5183, -0.6447],
          [ 1.8009,  1.8468, -0.9832, -1.5272]])
])
>>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0)
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  0.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  0.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  0.0000]]])
>>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6))
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276,  1.0000],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  1.0000,  1.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  1.0000,  1.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]]])
>>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2))
RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.

torch.nested.masked_select(tensor, mask) [исходный код]

Создаёт вложенный тензор на основе входного тензора со страйдовой компоновкой и маски со страйдовой компоновкой. В результирующем вложенном тензоре с рваной компоновкой сохраняются значения в тех позициях, где маска равна True. Размерность маски сохраняется и представляется с помощью offsets; этим функция отличается от masked_select(), в которой выходные данные сводятся к одномерному тензору.

Аргументы: tensor (torch.Tensor): тензор со страйдовой компоновкой, на основе которого создаётся вложенный тензор с рваной компоновкой. mask (torch.Tensor): тензор-маска со страйдовой компоновкой, применяемый к входному тензору

Пример:

>>> tensor = torch.randn(3, 3)
>>> mask = torch.tensor([[False, False, True], [True, False, True], [False, False, True]])
>>> nt = torch.nested.masked_select(tensor, mask)
>>> nt.shape  # xdoctest: +ELLIPSIS
torch.Size([3, j...])
>>> # Length of each item in the batch:
>>> nt.offsets().diff()
tensor([1, 2, 1])

>>> tensor = torch.randn(6, 5)
>>> mask = torch.tensor([False])
>>> nt = torch.nested.masked_select(tensor, mask)
>>> nt.shape  # xdoctest: +ELLIPSIS
torch.Size([6, j...])
>>> # Length of each item in the batch:
>>> nt.offsets().diff()
tensor([0, 0, 0, 0, 0, 0])

Тип возвращаемого значения:

Tensor

torch.nested.narrow(tensor, dim, start, length, layout=torch.strided) [исходный код]

Создаёт вложенный тензор (который может быть представлением) из tensor, тензора со страйдовой компоновкой. Функция имеет семантику, аналогичную torch.Tensor.narrow: в измерении dim новый вложенный тензор содержит только элементы из интервала [start, start+length). Поскольку вложенные представления допускают разные значения start и length в каждой «строке» этого измерения, start и length также могут быть тензорами формы tensor.shape[0].

Поведение зависит от выбранной компоновки вложенного тензора. При использовании страйдовой компоновки torch.narrow копирует суженные данные в непрерывный NT со страйдовой компоновкой, тогда как narrow() с рваной компоновкой создаёт неконтинуальное представление исходного тензора со страйдовой компоновкой. Такое представление очень удобно для хранения kv-кэшей в моделях Transformer: специализированные ядра SDPA могут легко работать с этим форматом, что повышает производительность.

Параметры:

• tensor (torch.Tensor) – тензор со страйдовой компоновкой, который будет использоваться как базовые данные вложенного тензора при рваной компоновке или копироваться при страйдовой компоновке.
• dim (int) – измерение, к которому применяется narrow. Для рваной компоновки поддерживается только dim=1, для страйдовой поддерживаются все измерения
• start (Union[int, torch.Tensor]) – начальный элемент для операции сужения
• length (Union[int, torch.Tensor]) – количество элементов, выбираемых при операции сужения

Именованные аргументы:

layout (torch.layout, необязательно) – требуемая компоновка возвращаемого вложенного тензора. Поддерживаются только страйдовая и рваная компоновки. По умолчанию: если None, страйдовая компоновка.

Тип возвращаемого значения:

Tensor

Пример:

>>> starts = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64)
>>> lengths = torch.tensor([3, 2, 2, 1, 5], dtype=torch.int64)
>>> narrow_base = torch.randn(5, 10, 20)
>>> nt_narrowed = torch.nested.narrow(narrow_base, 1, starts, lengths, layout=torch.jagged)
>>> nt_narrowed.is_contiguous()
False