Skip to content

Commit 7f7455a

Browse files
njzjznjzjz
committed
Avoid patching TensorFlow array objects
1 parent edf9333 commit 7f7455a

3 files changed

Lines changed: 109 additions & 216 deletions

File tree

‎src/array_api_compat/tensorflow/__init__.py‎

Lines changed: 2 additions & 210 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,4 @@
11
from collections.abc import Sequence
2-
from builtins import all as py_all
3-
from builtins import any as py_any
4-
from builtins import bool as py_bool
5-
from builtins import int as py_int
6-
from builtins import slice as py_slice
7-
from builtins import tuple as py_tuple
82
from typing import Final
93

104
import tensorflow as _tf
@@ -13,217 +7,15 @@
137

148
__all__ = clone_module("tensorflow", globals())
159

10+
# TensorShape is not the array object; this lets shape helpers consume x.shape
11+
# without mutating TensorFlow tensors.
1612
Sequence.register(type(_tf.TensorShape(())))
1713

18-
19-
class _ArrayAPIShape(py_tuple):
20-
def __new__(cls, tensor_shape):
21-
obj = super().__new__(cls, py_tuple(tensor_shape.as_list()))
22-
obj._tensor_shape = tensor_shape
23-
return obj
24-
25-
@property
26-
def rank(self):
27-
return self._tensor_shape.rank
28-
29-
@property
30-
def ndims(self):
31-
return self._tensor_shape.ndims
32-
33-
def as_list(self):
34-
return self._tensor_shape.as_list()
35-
36-
def __getitem__(self, key):
37-
out = self._tensor_shape[key]
38-
if isinstance(key, py_slice):
39-
return type(self)(out)
40-
return out
41-
42-
def __getattr__(self, name):
43-
return getattr(self._tensor_shape, name)
44-
45-
46-
def _patch_eager_tensor_shape() -> None:
47-
eager_tensor = type(_tf.constant(0))
48-
if getattr(eager_tensor, "_array_api_compat_shape_patched", False):
49-
return
50-
old_shape = eager_tensor.shape
51-
52-
def shape(self):
53-
return _ArrayAPIShape(old_shape.__get__(self, eager_tensor))
54-
55-
eager_tensor.shape = property(shape)
56-
eager_tensor._array_api_compat_shape_patched = True
57-
58-
59-
_patch_eager_tensor_shape()
60-
61-
62-
def _patch_eager_tensor_getitem() -> None:
63-
eager_tensor = type(_tf.constant(0))
64-
if getattr(eager_tensor, "_array_api_compat_getitem_patched", False):
65-
return
66-
old_getitem = eager_tensor.__getitem__
67-
68-
def tensor_index(key):
69-
if isinstance(key, _tf.Tensor):
70-
return int(key) if key.shape == () else key
71-
if isinstance(key, py_slice):
72-
return py_slice(tensor_index(key.start), tensor_index(key.stop), tensor_index(key.step))
73-
if isinstance(key, py_tuple):
74-
return py_tuple(tensor_index(k) for k in key)
75-
return key
76-
77-
def bool_getitem(self, key):
78-
x_shape = py_tuple(self.shape)
79-
key_shape = py_tuple(key.shape)
80-
if key.ndim > self.ndim or py_any(
81-
ks not in (xs, 0) for xs, ks in zip(x_shape, key_shape, strict=False)
82-
):
83-
raise IndexError("boolean index did not match indexed array")
84-
if key.ndim == 0:
85-
x = _tf.reshape(self, (1,) + x_shape)
86-
return _tf.boolean_mask(x, _tf.reshape(key, (1,)))
87-
if py_any(ks == 0 for ks in key_shape):
88-
return _tf.zeros((0,) + x_shape[key.ndim :], dtype=self.dtype)
89-
return _tf.boolean_mask(self, key)
90-
91-
def high_dim_int_getitem(self, key):
92-
shape = py_tuple(self.shape)
93-
if len(shape) <= 7:
94-
return None
95-
key = tensor_index(key)
96-
if isinstance(key, py_int):
97-
key = (key,)
98-
if not isinstance(key, py_tuple) or len(key) != len(shape):
99-
return None
100-
if not py_all(isinstance(k, py_int) and not isinstance(k, py_bool) for k in key):
101-
return None
102-
normalized = []
103-
for k, side in zip(key, shape, strict=True):
104-
k = k + side if k < 0 else k
105-
if k < 0 or k >= side:
106-
raise IndexError("index out of bounds")
107-
normalized.append(k)
108-
flat_index = 0
109-
for k, stride in zip(normalized, _strides(shape), strict=True):
110-
flat_index += k * stride
111-
return _tf.gather(_tf.reshape(self, (-1,)), flat_index)
112-
113-
def _strides(shape):
114-
strides = []
115-
stride = 1
116-
for side in reversed(shape):
117-
strides.append(stride)
118-
stride *= side
119-
return py_tuple(reversed(strides))
120-
121-
def getitem(self, key):
122-
if isinstance(key, _tf.Tensor) and key.dtype == _tf.bool:
123-
return bool_getitem(self, key)
124-
out = high_dim_int_getitem(self, key)
125-
if out is not None:
126-
return out
127-
return old_getitem(self, tensor_index(key))
128-
129-
eager_tensor.__getitem__ = getitem
130-
eager_tensor._array_api_compat_getitem_patched = True
131-
132-
133-
_patch_eager_tensor_getitem()
134-
135-
136-
def _patch_eager_tensor_binary_ops() -> None:
137-
eager_tensor = type(_tf.constant(0))
138-
if getattr(eager_tensor, "_array_api_compat_binary_ops_patched", False):
139-
return
140-
141-
def wrap(name, func_name, *, reverse=False):
142-
func = globals()[func_name]
143-
144-
def op(self, other):
145-
if reverse:
146-
return func(other, self)
147-
return func(self, other)
148-
149-
setattr(eager_tensor, name, op)
150-
151-
def wrap_unary(name, func_name):
152-
func = globals()[func_name]
153-
154-
def op(self):
155-
return func(self)
156-
157-
setattr(eager_tensor, name, op)
158-
159-
for name, func_name, reverse in [
160-
("__add__", "add", False),
161-
("__radd__", "add", True),
162-
("__sub__", "subtract", False),
163-
("__rsub__", "subtract", True),
164-
("__mul__", "multiply", False),
165-
("__rmul__", "multiply", True),
166-
("__floordiv__", "floor_divide", False),
167-
("__rfloordiv__", "floor_divide", True),
168-
("__mod__", "remainder", False),
169-
("__rmod__", "remainder", True),
170-
("__pow__", "pow", False),
171-
("__rpow__", "pow", True),
172-
("__lt__", "less", False),
173-
("__le__", "less_equal", False),
174-
("__gt__", "greater", False),
175-
("__ge__", "greater_equal", False),
176-
("__eq__", "equal", False),
177-
("__ne__", "not_equal", False),
178-
("__and__", "bitwise_and", False),
179-
("__rand__", "bitwise_and", True),
180-
("__or__", "bitwise_or", False),
181-
("__ror__", "bitwise_or", True),
182-
("__xor__", "bitwise_xor", False),
183-
("__rxor__", "bitwise_xor", True),
184-
("__lshift__", "bitwise_left_shift", False),
185-
("__rlshift__", "bitwise_left_shift", True),
186-
("__rshift__", "bitwise_right_shift", False),
187-
("__rrshift__", "bitwise_right_shift", True),
188-
("__matmul__", "matmul", False),
189-
("__rmatmul__", "matmul", True),
190-
]:
191-
wrap(name, func_name, reverse=reverse)
192-
193-
for name, func_name in [
194-
("__abs__", "abs"),
195-
("__invert__", "bitwise_invert"),
196-
("__neg__", "negative"),
197-
("__pos__", "positive"),
198-
]:
199-
wrap_unary(name, func_name)
200-
201-
eager_tensor._array_api_compat_binary_ops_patched = True
202-
203-
204-
def _patch_eager_tensor_dlpack() -> None:
205-
eager_tensor = type(_tf.constant(0))
206-
if getattr(eager_tensor, "_array_api_compat_dlpack_patched", False):
207-
return
208-
old_dlpack = eager_tensor.__dlpack__
209-
210-
def dlpack(self, *, stream=None, max_version=None, dl_device=None, copy=None):
211-
tensor = _tf.identity(self) if copy is True else self
212-
del dl_device
213-
return old_dlpack(tensor, stream=stream, max_version=max_version)
214-
215-
eager_tensor.__dlpack__ = dlpack
216-
eager_tensor._array_api_compat_dlpack_patched = True
217-
218-
21914
# These imports may overwrite names from the import * above.
22015
from . import _aliases
22116
from ._aliases import * # noqa: F403
22217
from ._info import __array_namespace_info__ # noqa: F401
22318

224-
_patch_eager_tensor_binary_ops()
225-
_patch_eager_tensor_dlpack()
226-
22719
__import__(__spec__.parent + ".linalg")
22820
__import__(__spec__.parent + ".fft")
22921

‎src/array_api_compat/tensorflow/linalg.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -210,8 +210,9 @@ def vector_norm(
210210
keepdims: bool = False,
211211
ord: JustInt | JustFloat = 2,
212212
) -> Array:
213+
out_dtype = _real_dtype_for(x.dtype)
213214
if axis == ():
214-
return tf.cast(x != 0, x.dtype) if ord == 0 else abs(x)
215+
return tf.cast(x != 0, out_dtype) if ord == 0 else tf.cast(abs(x), out_dtype)
215216

216217
if axis is None:
217218
x_ = tf.reshape(x, (-1,))
@@ -230,7 +231,6 @@ def vector_norm(
230231
axis_ = axis
231232

232233
abs_x = tf.abs(x_)
233-
out_dtype = _real_dtype_for(x.dtype)
234234
abs_x = tf.cast(abs_x, out_dtype)
235235
if ord == 0:
236236
out = tf.cast(count_nonzero(x_, axis=axis_), out_dtype)

‎tensorflow-skips.txt‎

Lines changed: 105 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,125 @@
1-
# We do not wrap the TensorFlow Tensor object. TensorFlow tensors are immutable
2-
# and do not implement the full Array API array object surface.
1+
# We do not wrap or patch the TensorFlow Tensor object. TensorFlow tensors are
2+
# immutable and do not implement the full Array API array object surface. Tests
3+
# below exercise tf.Tensor itself, or use tf.Tensor.shape/dunder operators while
4+
# constructing expected results.
35
array_api_tests/test_creation_functions.py::test_asarray_arrays
46

57
array_api_tests/test_array_object.py::test_setitem
68
array_api_tests/test_array_object.py::test_setitem_masking
9+
array_api_tests/test_array_object.py::test_getitem_masking
710
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1
811
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2
912

10-
# TensorFlow's native true-division dunder still handles array-array operands
11-
# before our namespace-level dtype promotion can run.
13+
# TensorFlow exposes shape as TensorShape, not tuple. These tests feed x.shape
14+
# back into Hypothesis strategies which require tuple shapes.
15+
array_api_tests/test_data_type_functions.py::test_broadcast_to
16+
array_api_tests/test_indexing_functions.py::test_take_along_axis
17+
array_api_tests/test_manipulation_functions.py::test_concat
18+
array_api_tests/test_manipulation_functions.py::test_reshape
19+
array_api_tests/test_manipulation_functions.py::test_stack
20+
21+
# These tests validate results using tf.Tensor indexing/dunders directly.
22+
array_api_tests/test_dlpack.py::test_dunder_dlpack
23+
array_api_tests/test_linalg.py::test_cross
24+
array_api_tests/test_linalg.py::test_outer
25+
array_api_tests/test_linalg.py::test_solve
26+
array_api_tests/test_manipulation_functions.py::test_repeat
27+
array_api_tests/test_operators_and_elementwise_functions.py::test_clip
28+
array_api_tests/test_operators_and_elementwise_functions.py::test_where_with_scalars
29+
array_api_tests/test_searching_functions.py::test_where
30+
31+
# TensorFlow's native dunder operators handle operands before our
32+
# namespace-level dtype promotion can run.
33+
array_api_tests/test_operators_and_elementwise_functions.py::test_abs[__abs__]
34+
array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x1, x2)]
35+
array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x, s)]
36+
array_api_tests/test_operators_and_elementwise_functions.py::test_add[__iadd__(x1, x2)]
37+
array_api_tests/test_operators_and_elementwise_functions.py::test_add[__iadd__(x, s)]
38+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__and__(x1, x2)]
39+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__and__(x, s)]
40+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__iand__(x1, x2)]
41+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_and[__iand__(x, s)]
42+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x1, x2)]
43+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x, s)]
44+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x1, x2)]
45+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x, s)]
46+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_invert[__invert__]
47+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__or__(x1, x2)]
48+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__or__(x, s)]
49+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__ior__(x1, x2)]
50+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_or[__ior__(x, s)]
51+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x1, x2)]
52+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x, s)]
53+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x1, x2)]
54+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x, s)]
55+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__xor__(x1, x2)]
56+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__xor__(x, s)]
57+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__ixor__(x1, x2)]
58+
array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_xor[__ixor__(x, s)]
1259
array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__truediv__(x1, x2)]
60+
array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__truediv__(x, s)]
1361
array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__itruediv__(x1, x2)]
62+
array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__itruediv__(x, s)]
63+
array_api_tests/test_operators_and_elementwise_functions.py::test_equal[__eq__(x1, x2)]
64+
array_api_tests/test_operators_and_elementwise_functions.py::test_equal[__eq__(x, s)]
65+
array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__floordiv__(x1, x2)]
66+
array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__floordiv__(x, s)]
67+
array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__ifloordiv__(x1, x2)]
68+
array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__ifloordiv__(x, s)]
69+
array_api_tests/test_operators_and_elementwise_functions.py::test_greater[__gt__(x1, x2)]
70+
array_api_tests/test_operators_and_elementwise_functions.py::test_greater[__gt__(x, s)]
71+
array_api_tests/test_operators_and_elementwise_functions.py::test_greater_equal[__ge__(x1, x2)]
72+
array_api_tests/test_operators_and_elementwise_functions.py::test_greater_equal[__ge__(x, s)]
73+
array_api_tests/test_operators_and_elementwise_functions.py::test_less[__lt__(x1, x2)]
74+
array_api_tests/test_operators_and_elementwise_functions.py::test_less[__lt__(x, s)]
75+
array_api_tests/test_operators_and_elementwise_functions.py::test_less_equal[__le__(x1, x2)]
76+
array_api_tests/test_operators_and_elementwise_functions.py::test_less_equal[__le__(x, s)]
77+
array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__mul__(x1, x2)]
78+
array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__mul__(x, s)]
79+
array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__imul__(x1, x2)]
80+
array_api_tests/test_operators_and_elementwise_functions.py::test_multiply[__imul__(x, s)]
81+
array_api_tests/test_operators_and_elementwise_functions.py::test_negative[__neg__]
82+
array_api_tests/test_operators_and_elementwise_functions.py::test_not_equal[__ne__(x1, x2)]
83+
array_api_tests/test_operators_and_elementwise_functions.py::test_not_equal[__ne__(x, s)]
84+
array_api_tests/test_operators_and_elementwise_functions.py::test_positive[__pos__]
85+
array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__pow__(x1, x2)]
86+
array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__pow__(x, s)]
87+
array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__ipow__(x1, x2)]
88+
array_api_tests/test_operators_and_elementwise_functions.py::test_pow[__ipow__(x, s)]
89+
array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__mod__(x1, x2)]
90+
array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__mod__(x, s)]
91+
array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__imod__(x1, x2)]
92+
array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__imod__(x, s)]
93+
array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub__(x1, x2)]
94+
array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub__(x, s)]
95+
array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x1, x2)]
96+
array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__isub__(x, s)]
97+
98+
array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is -0 and x2_i > 0) -> +0]
99+
array_api_tests/test_special_cases.py::test_binary[__mod__(x1_i is +0 and x2_i < 0) -> -0]
100+
array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is -0 and x2_i > 0) -> +0]
101+
array_api_tests/test_special_cases.py::test_iop[__imod__(x1_i is +0 and x2_i < 0) -> -0]
102+
array_api_tests/test_special_cases.py::test_nan_propagation[max]
103+
array_api_tests/test_special_cases.py::test_nan_propagation[mean]
104+
array_api_tests/test_special_cases.py::test_nan_propagation[min]
105+
array_api_tests/test_special_cases.py::test_nan_propagation[prod]
106+
array_api_tests/test_special_cases.py::test_nan_propagation[std]
107+
array_api_tests/test_special_cases.py::test_nan_propagation[sum]
108+
array_api_tests/test_special_cases.py::test_nan_propagation[var]
14109

15110
array_api_tests/test_has_names.py::test_has_names[array_method-__array_namespace__]
111+
array_api_tests/test_has_names.py::test_has_names[array_method-__lshift__]
112+
array_api_tests/test_has_names.py::test_has_names[array_method-__pos__]
113+
array_api_tests/test_has_names.py::test_has_names[array_method-__rshift__]
16114
array_api_tests/test_has_names.py::test_has_names[array_method-__setitem__]
17115
array_api_tests/test_has_names.py::test_has_names[array_method-to_device]
18116
array_api_tests/test_has_names.py::test_has_names[array_attribute-T]
19117
array_api_tests/test_has_names.py::test_has_names[array_attribute-mT]
20118
array_api_tests/test_has_names.py::test_has_names[array_attribute-size]
21119

22120
array_api_tests/test_signatures.py::test_array_method_signature[__array_namespace__]
121+
array_api_tests/test_signatures.py::test_array_method_signature[__lshift__]
122+
array_api_tests/test_signatures.py::test_array_method_signature[__pos__]
123+
array_api_tests/test_signatures.py::test_array_method_signature[__rshift__]
23124
array_api_tests/test_signatures.py::test_array_method_signature[__setitem__]
24125
array_api_tests/test_signatures.py::test_array_method_signature[to_device]

0 commit comments

Comments
 (0)