Skip to content

merge_axes

Merge multiple tensor axes into a single axis using isometry fusion.

merge_axes 

merge_axes(
    tensor: Tensor,
    axes: Sequence[Union[int, str]],
    *,
    merged_tag: Optional[str] = None,
    direction: Direction = Direction.OUT,
) -> Tuple[Tensor, Tensor]

Merge multiple tensor axes into a single axis using isometry fusion.

This function creates an n-to-1 isometry that fuses the specified axes, contracts it with the input tensor to perform the merging, and returns both the merged tensor and the conjugate of the isometry (which can be used to unfuse the axis later).

Parameters:

Name Type Description Default
tensor Tensor

Input tensor whose axes should be merged.

 required
axes Sequence[Union[int, str]]

Sequence of axes to merge. Each element can be either an integer position (0-indexed) or a string itag. Must specify at least 2 axes.

 required
merged_tag Optional[str]

Optional tag for the merged axis in the result. If None, uses "merged".

 None
direction Direction

Direction for the merged axis. Defaults to Direction.OUT.

 OUT

Returns:

Name Type Description
merged_tensor Tensor

Tensor with the specified axes merged into a single axis. The merged axis appears first, followed by the remaining unmerged axes in their original order.
isometry_conj Tensor

Conjugate of the isometry used for merging. Can be contracted with the merged tensor to unfuse the axis back to the original indices.

Raises:

Type Description
ValueError:

If fewer than 2 axes are specified, if axis specifications are invalid, or if axes don't exist in the tensor.
TypeError:

If axis specifications are neither int nor str.

Examples:

Merge three axes of a tensor:

>>> tensor = Tensor.random([idx1, idx2, idx3, idx4], itags=['a', 'b', 'c', 'd'])
>>> merged, iso_conj = merge_axes(tensor, ['a', 'b', 'c'], merged_tag='abc')
>>> merged.itags  # Merged index 'abc' appears first
('abc', 'd')

Merge using integer positions:

>>> merged, iso_conj = merge_axes(tensor, [0, 1, 2], merged_tag='merged')

Unfuse the merged axis (approximately recover original):

>>> from nicole import contract
>>> unmerged = contract(merged, iso_conj)  # Should match original structure
Notes
  • The merged axis will have direction opposite to what natural fusion produces unless specified otherwise via the direction parameter
  • The isometry conjugate has all its indices flipped, making it suitable for contracting with the merged tensor to reverse the operation
  • Axes are merged in the order they appear in the tensor, not the order specified in the axes parameter

Description

Creates an n-to-1 isometry that fuses the specified axes, contracts it with the input tensor to perform the merging, and returns both the merged tensor and the conjugate of the isometry (which can be used to unfuse the axis later).

This is useful for reducing tensor order while preserving all information, with the ability to reverse the operation.

Usage

Merge by tags

from nicole import Tensor, merge_axes

# Create a 4-index tensor
tensor = Tensor.random([idx1, idx2, idx3, idx4], itags=['a', 'b', 'c', 'd'])

# Merge three axes
merged, iso_conj = merge_axes(tensor, ['a', 'b', 'c'], merged_tag='abc')
print(merged.itags)  # ('abc', 'd')

Merge by positions

# Merge using integer positions (0-indexed)
merged, iso_conj = merge_axes(tensor, [0, 1, 2], merged_tag='merged')

Unfuse merged axis

from nicole import contract

# Reverse the merging operation
unmerged = contract(merged, iso_conj)
# unmerged should approximately match the original tensor structure

See Also

Notes

  • Must specify at least 2 axes to merge
  • Axes can be specified as integer positions (0-indexed) or string tags
  • The merged axis appears first in the output tensor
  • The returned isometry conjugate can contract with the merged tensor to reverse the operation
  • Axes are merged in the order they appear in the tensor, not the order specified
  • The merged axis direction can be controlled via the direction parameter