Repository navigation
Expand file tree
/
Copy pathvalidation.py
More file actions
141 lines (117 loc) · 4.83 KB
/
Copy pathvalidation.py
File metadata and controls
141 lines (117 loc) · 4.83 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
# ===----------------------------------------------------------------------=== #
# 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.
# ===----------------------------------------------------------------------=== #
"""Helper functions for asserting in ops."""
import warnings
from collections.abc import Iterable
from ..graph import DevicePlacementPolicy, Graph
from ..value import BufferValue, TensorValue
def _check_device_placement(op_name: str, todo: str = "") -> None:
"""Enforce the active graph's device placement policy for CPU-only ops.
Call this before implicitly transferring a non-CPU tensor to CPU inside an
op implementation. Follow it immediately with ``transfer_to(x, CPU())``.
Behavior is driven by ``Graph.strict_device_placement``:
- ``Ignore``: no action.
- ``Warn`` (default): emit a ``UserWarning``.
- ``Error``: raise ``ValueError``.
Args:
op_name: The public op name, for example, ``"ops.nonzero"``.
todo: Tracking ticket for GPU kernel work, for example, ``"TODO(GEX-2041)."``.
"""
policy = Graph.current.strict_device_placement
if policy is DevicePlacementPolicy.Ignore:
return
msg = (
f"`{op_name}` does not support GPU execution and will implicitly"
" transfer the tensor to CPU."
)
if todo:
msg += f" {todo}"
msg += (
" Call `ops.transfer_to(x, CPU())` explicitly, or construct the"
" Graph with `strict_device_placement=DevicePlacementPolicy.Error`"
" to make this an error."
)
if policy is DevicePlacementPolicy.Error:
raise ValueError(msg)
warnings.warn(msg, stacklevel=3)
def indent(
lines: Iterable[str], level: int = 1, indent: str = " "
) -> Iterable[str]:
"""Prepends each line with the given indent string repeated level times."""
for line in lines:
yield (indent * level) + line
def assert_same_device(
*values: TensorValue | BufferValue,
**named_values: TensorValue | BufferValue,
) -> None:
"""Raises ValueError if any of the given values are not on the same device."""
named_values = {
**{str(i): value for i, value in enumerate(values)},
**named_values,
}
if len({v.device for v in named_values.values()}) > 1:
raise ValueError(
"Input values must be on the same device\n"
+ "\n".join(
indent(
f"{name}: {value.type}"
for name, value in named_values.items()
)
)
)
def assert_on_host(
*values: TensorValue | BufferValue,
**named_values: TensorValue | BufferValue,
) -> None:
"""Raises ValueError if any of the given values are not on the CPU device."""
named_values = {
**{str(i): value for i, value in enumerate(values)},
**named_values,
}
not_on_host = {
k: v for k, v in named_values.items() if not v.device.is_cpu()
}
if not_on_host:
raise ValueError(
"Input value must be on the host device\n"
+ "\n".join(
indent(
f"{name}: {value.type}"
for name, value in not_on_host.items()
)
)
)
def assert_same_shape(*values: TensorValue | BufferValue) -> None:
"""Raises ValueError if the given values do not all have the same shape."""
first_shape = values[0].shape
for i, tensor in enumerate(values[1:], start=1):
if tensor.shape != first_shape:
raise ValueError(
f"All input tensors must have the same shape. "
f"Input 0 has shape {first_shape}, but input {i} has shape {tensor.shape}"
)
def assert_same_dtype(*values: TensorValue | BufferValue) -> None:
"""Raises ValueError if the given values do not all have the same dtype."""
first_dtype = values[0].dtype
for i, tensor in enumerate(values[1:], start=1):
if tensor.dtype != first_dtype:
raise ValueError(
f"All input tensors must have the same dtype. "
f"Input 0 has dtype {first_dtype}, but input {i} has dtype {tensor.dtype}"
)
def assert_valid_axis(value: TensorValue | BufferValue, axis: int) -> None:
"""Raises IndexError if axis is out of range for the value's rank."""
if not (-value.rank <= axis < value.rank):
raise IndexError(
f"Axis must be in range [-{value.rank}, {value.rank}), got {axis}"
)