Repository navigation
Expand file tree
/
Copy pathmatmul.py
More file actions
111 lines (93 loc) · 4.16 KB
/
Copy pathmatmul.py
File metadata and controls
111 lines (93 loc) · 4.16 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
102
103
104
105
106
107
108
109
110
111
# ===----------------------------------------------------------------------=== #
# 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 matmul."""
from max._core.dialects import rmo
from max.dtype import DType
from ..graph import Graph
from ..type import TensorType
from ..value import TensorValue, TensorValueLike
from .validation import assert_same_device
def matmul(
lhs: TensorValueLike,
rhs: TensorValueLike,
*,
out_dtype: DType | None = None,
) -> TensorValue:
"""Computes the matrix product of two tensors.
Use matrix multiplication to implement key building blocks like linear
transformations, attention mechanisms, and fully connected layers. You can
call ``matmul()`` directly or use the ``@`` operator, which calls
``matmul()`` implicitly.
.. 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("matmul_example") as graph:
lhs = ops.constant(
[[1.0, 2.0], [3.0, 4.0]], DType.float32, device=device
)
rhs = ops.constant(
[[5.0, 6.0], [7.0, 8.0]], DType.float32, device=device
)
graph.output(ops.matmul(lhs, rhs, out_dtype=DType.float32))
model = InferenceSession().load(graph)
result = model.execute()[0]
.. invisible-code-block: python
import numpy as np
assert np.allclose(result.to_numpy(), [[19.0, 22.0], [43.0, 50.0]])
The innermost two dimensions of each input are treated as a matrix.
In the example above ``lhs`` has shape ``(M, K) = (2, 2)`` and ``rhs``
has shape ``(K, N) = (2, 2)``, producing an output of shape
``(M, N) = (2, 2)``. The ``K`` dimensions must match. Any remaining
outer (batch) dimensions are broadcast.
If ``lhs`` is 1-D it is reshaped to ``1xD``, and if ``rhs`` is 1-D it is
reshaped to ``Dx1``. In both cases, the added size-1 dimensions are
removed from the output shape.
The result has dtype ``out_dtype``. The inputs keep their own dtypes and
are never cast, so ``bfloat16`` inputs with ``out_dtype=DType.float32``
keep the float32 accumulator instead of rounding it to bfloat16.
Args:
lhs: The left-hand side input tensor.
rhs: The right-hand side input tensor.
out_dtype: The dtype of the result. Defaults to the inputs' dtype,
which they must then share.
Returns:
A tensor value representing the matrix product of ``lhs`` and ``rhs``.
For 2-D inputs, the output shape is ``(M, N)`` where ``lhs`` is
``(M, K)`` and ``rhs`` is ``(K, N)``. For higher-dimensional inputs,
batch dimensions are preserved and the operation is applied to the
last two dimensions of each input.
Raises:
ValueError: If ``out_dtype`` is not given and ``lhs`` and ``rhs``
have different dtypes.
"""
lhs = TensorValue(lhs)
rhs = TensorValue(rhs)
assert_same_device(lhs=lhs, rhs=rhs)
if out_dtype is None:
if lhs.dtype != rhs.dtype:
raise ValueError(
f"matmul inputs have different dtypes ({lhs.dtype} and"
f" {rhs.dtype}); pass out_dtype to choose the result dtype"
)
out_dtype = lhs.dtype
inferred = Graph.current._add_op_generated(
rmo.MatmulOp, input_x=lhs, input_y=rhs
)[0].tensor
# rmo.matmul infers only the shape; the result dtype is the caller's.
result = inferred.to_mlir()
result.type = TensorType(
out_dtype, inferred.shape, inferred.device
).to_mlir()
return TensorValue(result)