Repository navigation
Expand file tree
/
Copy pathtranspose.py
More file actions
101 lines (82 loc) · 3.36 KB
/
Copy pathtranspose.py
File metadata and controls
101 lines (82 loc) · 3.36 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# ===----------------------------------------------------------------------=== #
# Copyright (c) 2026, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
"""Op implementation for transpose."""
from max._core.dialects import kgen, rmo
from max.dtype import DType
from ..graph import Graph
from ..type import DeviceRef
from ..value import TensorType, TensorValue, TensorValueLike
from .constant import constant
from .utils import check_axis_in_bounds
def _axis_bounds(rank: int) -> tuple[int, int]:
if rank == 0:
return -1, 0
return -rank, rank - 1
def transpose(x: TensorValueLike, axis_1: int, axis_2: int) -> TensorValue:
"""Transposes two axes of a symbolic tensor.
.. code-block:: python
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops
device = DeviceRef.CPU()
with Graph("transpose") as graph:
# x has shape (2, 3).
x = ops.constant(
[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
DType.float32,
device=device,
)
# Swap axes 0 and 1, producing shape (3, 2).
graph.output(ops.transpose(x, 0, 1))
model = InferenceSession().load(graph)
result = model.execute()[0]
Args:
x: The input symbolic tensor to transpose.
axis_1: One of the two axes to transpose. If negative, this indexes from the
end of the tensor. For example, a value of ``-1`` refers to the last
axis.
axis_2: The other axis to transpose. If negative, this indexes from the end
of the tensor.
Returns:
A ``TensorValue`` representing the input tensor with ``axis_1`` and
``axis_2`` transposed. It has the same elements and dtype as ``x``,
with the order of the elements changed according to the transposition.
For a rank-zero tensor, axes ``-1`` and ``0`` are accepted and the
scalar is returned unchanged.
Raises:
IndexError: If ``axis_1`` or ``axis_2`` is out of range.
"""
v = TensorValue(x)
rank = len(v.shape)
check_axis_in_bounds(axis_1, rank, _axis_bounds, "axis_1")
check_axis_in_bounds(axis_2, rank, _axis_bounds, "axis_2")
if axis_1 < 0:
axis_1 += rank
if axis_2 < 0:
axis_2 += rank
new_shape = v.shape
indices = list(range(len(new_shape)))
# Only change the shape for non-zero rank tensors.
if rank > 0:
new_shape[axis_1], new_shape[axis_2] = (
new_shape[axis_2],
new_shape[axis_1],
)
indices[axis_1], indices[axis_2] = axis_2, axis_1
return Graph.current._add_op_generated(
rmo.MoTransposeOp,
TensorType(dtype=v.dtype, shape=new_shape, device=v.device),
v,
constant(indices, DType.int64, DeviceRef.CPU()),
kgen.ParamDeclArrayAttr([]),
)[0].tensor