-
Notifications
You must be signed in to change notification settings - Fork 16
Expand file tree
/
Copy pathselect.py
More file actions
303 lines (263 loc) · 12.2 KB
/
Copy pathselect.py
File metadata and controls
303 lines (263 loc) · 12.2 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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
import inspect
from ... import select
from ...dtypes import BOOL, UINT64
from .. import _has_numba
from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op
from .indexunary import IndexUnaryOp, TypedBuiltinIndexUnaryOp
if _has_numba:
from .base import _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper
class TypedBuiltinSelectOp(TypedOpBase):
__slots__ = ()
opclass = "SelectOp"
def __call__(self, val, thunk=None):
if thunk is None:
thunk = False # most basic form of 0 when unifying dtypes
return _call_op(self, val, thunk=thunk)
thunk_type = TypedBuiltinIndexUnaryOp.thunk_type
class TypedUserSelectOp(TypedOpBase):
__slots__ = ()
opclass = "SelectOp"
_owns_gb_obj = True # underlying object is a GrB_IndexUnaryOp
def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None):
super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2)
@property
def orig_func(self):
return self.parent.orig_func
@property
def _numba_func(self):
return self.parent._numba_func
thunk_type = TypedBuiltinSelectOp.thunk_type
__call__ = TypedBuiltinSelectOp.__call__
class ParameterizedSelectOp(ParameterizedUdf):
__slots__ = "func", "__signature__", "_is_udt"
def __init__(self, name, func, *, anonymous=False, is_udt=False):
self.func = func
self.__signature__ = inspect.signature(func)
self._is_udt = is_udt
if name is None:
name = getattr(func, "__name__", name)
super().__init__(name, anonymous)
def _call(self, *args, **kwargs):
sel = self.func(*args, **kwargs)
sel._parameterized_info = (self, args, kwargs)
return SelectOp.register_anonymous(sel, self.name, is_udt=self._is_udt)
class SelectOp(OpBase):
"""Identical to an :class:`IndexUnaryOp <graphblas.core.operator.IndexUnaryOp>`,
but must have a Boolean return type.
A SelectOp is used exclusively to select a subset of values from a collection where
the function returns True.
Built-in and registered SelectOps are located in the ``graphblas.select`` namespace.
"""
__slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func"
_module = select
_modname = "select"
_custom_dtype = None
_typed_class = TypedBuiltinSelectOp
_typed_user_class = TypedUserSelectOp
@classmethod
def _from_indexunary(cls, iop):
obj = cls(
iop.name,
iop.orig_func,
anonymous=iop._anonymous,
is_positional=iop.is_positional,
is_udt=iop._is_udt,
numba_func=iop._numba_func,
)
if not all(x == BOOL for x in iop.types.values()):
raise ValueError("SelectOp must have BOOL return type")
for type_, t in iop._typed_ops.items():
if iop.orig_func is not None:
op = cls._typed_user_class(
obj,
iop.name,
t.type,
t.return_type,
t.gb_obj,
)
# Borrow the IndexUnaryOp's allocation instead of making a
# second one. Holding ``t`` keeps that handle alive for as long
# as this SelectOp can use it: ``iop`` is a temporary in
# ``register_anonymous``, so without this the handle is freed
# the moment it is collected and every call raises
# UninitializedObject.
op._gb_obj_owner = t
else:
op = cls._typed_class(
obj,
iop.name,
t.type,
t.return_type,
t.gb_obj,
t.gb_name,
)
# type is not always equal to t.type, so can't use op._add
# but otherwise perform the same logic
obj._typed_ops[type_] = op
obj.types[type_] = op.return_type
return obj
def _compile_udt(self, dtype, dtype2):
if dtype2 is None: # pragma: no cover
dtype2 = dtype
dtypes = (dtype, dtype2)
if dtypes in self._udt_types:
return self._udt_ops[dtypes]
if self._numba_func is None:
raise KeyError(f"{self.name} does not work with {dtypes} types")
# It would be nice if we could reuse compiling done for IndexUnaryOp
numba_func = self._numba_func
sig = (dtype.numba_type, UINT64.numba_type, UINT64.numba_type, dtype2.numba_type)
_compile_udf_for_udt(
numba_func, sig, op_kind="select", op_name=self.name, dtypes=(dtype, dtype2)
)
select_wrapper, wrapper_sig = _get_udt_wrapper(
numba_func, BOOL, dtype, dtype2, include_indexes=True
)
return _finalize_udt_op(
self, dtype, dtype2, BOOL, select_wrapper, wrapper_sig, TypedUserSelectOp
)
@classmethod
def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False):
"""Register a SelectOp without registering it in the ``graphblas.select`` namespace.
Because it is not registered in the namespace, the name is optional.
The return type must be Boolean.
Parameters
----------
func : FunctionType
The function to compile. For all current backends, this must be able
to be compiled with ``numba.njit``.
``func`` takes four input parameters (any dtype, int64, int64,
any dtype) and returns boolean. The first argument (any dtype) is
the value of the input Matrix or Vector, the second argument (int64)
is the row index of the Matrix or the index of the Vector, the third
argument (int64) is the column index of the Matrix or 0 for a Vector,
and the fourth argument (any dtype) is the value of the input Scalar.
name : str, optional
The name of the operator. This *does not* show up as ``gb.select.{name}``.
parameterized : bool, default False
When True, create a parameterized user-defined operator, which means
additional parameters can be "baked into" the operator when used.
For example, ``gb.binary.isclose`` is a parameterized BinaryOp that
optionally accepts ``rel_tol`` and ``abs_tol`` parameters, and it
can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``.
When creating a parameterized user-defined operator, the ``func``
parameter must be a callable that *returns* a function that will
then get compiled.
is_udt : bool, default False
Whether the operator is intended to operate on user-defined types.
If True, then the function will not be automatically compiled for
builtin types, and it will be compiled "just in time" when used.
Setting ``is_udt=True`` is also helpful when the left and right
dtypes need to be different.
Returns
-------
SelectOp or ParameterizedSelectOp
"""
cls._check_supports_udf("register_anonymous")
if parameterized:
return ParameterizedSelectOp(name, func, anonymous=True, is_udt=is_udt)
iop = IndexUnaryOp._build(name, func, anonymous=True, is_udt=is_udt)
return SelectOp._from_indexunary(iop)
@classmethod
def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False):
"""Register a new SelectOp and save it to ``graphblas.select`` namespace.
The function will also be registered as a IndexUnaryOp with the same name.
The return type must be Boolean.
Parameters
----------
name : str
The name of the operator. This will show up as ``gb.select.{name}``.
The name may contain periods, ".", which will result in nested objects
such as ``gb.select.x.y.z`` for name ``"x.y.z"``.
func : FunctionType
The function to compile. For all current backends, this must be able
to be compiled with ``numba.njit``.
``func`` takes four input parameters (any dtype, int64, int64,
any dtype) and returns boolean. The first argument (any dtype) is
the value of the input Matrix or Vector, the second argument (int64)
is the row index of the Matrix or the index of the Vector, the third
argument (int64) is the column index of the Matrix or 0 for a Vector,
and the fourth argument (any dtype) is the value of the input Scalar.
parameterized : bool, default False
When True, create a parameterized user-defined operator, which means
additional parameters can be "baked into" the operator when used.
For example, ``gb.binary.isclose`` is a parameterized BinaryOp that
optionally accepts ``rel_tol`` and ``abs_tol`` parameters, and it
can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``.
When creating a parameterized user-defined operator, the ``func``
parameter must be a callable that *returns* a function that will
then get compiled.
is_udt : bool, default False
Whether the operator is intended to operate on user-defined types.
If True, then the function will not be automatically compiled for
builtin types, and it will be compiled "just in time" when used.
Setting ``is_udt=True`` is also helpful when the left and right
dtypes need to be different.
lazy : bool, default False
If False (the default), then the function will be automatically
compiled for builtin types (unless ``is_udt`` is True).
Compiling functions can be slow, however, so you may want to
delay compilation and only compile when the operator is used,
which is done by setting ``lazy=True``.
Examples
--------
>>> gb.select.register_new("upper_left_triangle", lambda x, i, j, thunk: i + j <= thunk)
>>> dir(gb.select)
[..., 'upper_left_triangle', ...]
"""
cls._check_supports_udf("register_new")
iop = IndexUnaryOp.register_new(
name, func, parameterized=parameterized, is_udt=is_udt, lazy=lazy
)
module, funcname = cls._remove_nesting(name, strict=False)
if lazy:
module._delayed[funcname] = (
cls._get_delayed,
{"name": name},
)
elif parameterized:
op = ParameterizedSelectOp(funcname, func, is_udt=is_udt)
setattr(module, funcname, op)
return op
elif not all(x == BOOL for x in iop.types.values()):
# Undo registration of indexunaryop
imodule, funcname = IndexUnaryOp._remove_nesting(name, strict=False)
delattr(imodule, funcname)
raise ValueError("SelectOp must have BOOL return type")
else:
return getattr(module, funcname)
@classmethod
def _get_delayed(cls, name):
imodule, funcname = IndexUnaryOp._remove_nesting(name, strict=False)
iop = getattr(imodule, name)
if not all(x == BOOL for x in iop.types.values()):
raise ValueError("SelectOp must have BOOL return type")
module, funcname = cls._remove_nesting(name, strict=False)
return getattr(module, funcname)
@classmethod
def _initialize(cls):
if cls._initialized: # pragma: no cover (safety)
return
# IndexUnaryOp adds it boolean-returning objects to SelectOp
IndexUnaryOp._initialize()
cls._initialized = True
def __init__(
self,
name,
func=None,
*,
anonymous=False,
is_positional=False,
is_udt=False,
numba_func=None,
):
super().__init__(name, anonymous=anonymous)
self.orig_func = func
self._numba_func = numba_func
self.is_positional = is_positional
self._is_udt = is_udt
if is_udt:
self._udt_types = {} # {dtype: DataType}
self._udt_ops = {} # {dtype: TypedUserIndexUnaryOp}
__call__ = TypedBuiltinSelectOp.__call__
ParameterizedSelectOp._op_class = SelectOp