feat: implement array compatibility/interface protocols for virtual array and placeholder array (for better errors)#3839
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files
🚀 New features to boost your workflow:
|
|
The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3839 |
|
@pfackeldey we've been iterating over numpy arrays this whole time in certain edge cases. If you do |
ikrommyd
changed the title
|
This is actually way more critical than I thought. I added the following code to catch and raise an error when a module like numpy cupy or jax try to iterate over an array because they don't find the appopriate way to convert it. diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py
index 7f8da1ec..6eb45a5a 100644
--- a/src/awkward/_nplikes/typetracer.py
+++ b/src/awkward/_nplikes/typetracer.py
@@ -405,6 +405,32 @@ class TypeTracerArray(NDArrayOperatorsMixin, ArrayLike):
"Bug in Awkward Array: cannot convert TypeTracerArray into a concrete array"
)
+ _cycler = 0
+
+ def __getattr__(self, name):
+ if name == "__array_struct__":
+ self._cycler = 1
+ raise AttributeError(name)
+ elif name == "__array_interface__":
+ if self._cycler == 1:
+ self._cycler = 2
+ else:
+ self._cycler = 0
+ return self._data.__array_interface__
+ elif name == "__array__":
+ if self._cycler == 2:
+ self._cycler = 3
+ raise Exception(
+ "Module is coming to the conclusion that VirtualNDArray is a Python iterable"
+ )
+ else:
+ self._cycler = 0
+ raise AttributeError(name)
+ elif name in self.__dict__:
+ return self.__dict__[name]
+ else:
+ raise AttributeError(name)
+
class _CTypes:
data = 0
diff --git a/src/awkward/_nplikes/virtual.py b/src/awkward/_nplikes/virtual.py
index 0c517423..6eafe6db 100644
--- a/src/awkward/_nplikes/virtual.py
+++ b/src/awkward/_nplikes/virtual.py
@@ -413,3 +413,29 @@ class VirtualNDArray(NDArrayOperatorsMixin, MaterializableArray):
def __reduce__(self):
return self.materialize().__reduce__()
+
+ _cycler = 0
+
+ def __getattr__(self, name):
+ if name == "__array_struct__":
+ self._cycler = 1
+ raise AttributeError(name)
+ elif name == "__array_interface__":
+ if self._cycler == 1:
+ self._cycler = 2
+ else:
+ self._cycler = 0
+ return self._data.__array_interface__
+ elif name == "__array__":
+ if self._cycler == 2:
+ self._cycler = 3
+ raise Exception(
+ "Module is coming to the conclusion that VirtualNDArray is a Python iterable"
+ )
+ else:
+ self._cycler = 0
+ raise AttributeError(name)
+ elif name in self.__dict__:
+ return self.__dict__[name]
+ else:
+ raise AttributeError(name)
diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py
index 2148e31f..801149f1 100644
--- a/src/awkward/contents/content.py
+++ b/src/awkward/contents/content.py
@@ -1406,6 +1406,32 @@ class Content(Meta):
def _arrow_needs_option_type(cls):
return cls.is_option # is_option is a class property of Meta
+ _cycler = 0
+
+ def __getattr__(self, name):
+ if name == "__array_struct__":
+ self._cycler = 1
+ raise AttributeError(name)
+ elif name == "__array_interface__":
+ if self._cycler == 1:
+ self._cycler = 2
+ else:
+ self._cycler = 0
+ return self._data.__array_interface__
+ elif name == "__array__":
+ if self._cycler == 2:
+ self._cycler = 3
+ raise Exception(
+ "Module is coming to the conclusion that NumpyArray is a Python iterable"
+ )
+ else:
+ self._cycler = 0
+ raise AttributeError(name)
+ elif name in self.__dict__:
+ return self.__dict__[name]
+ else:
+ raise AttributeError(name)
+
@register_backend_lookup_factory
def find_content_backend(obj: type):
diff --git a/src/awkward/index.py b/src/awkward/index.py
index f709930d..1014c5e3 100644
--- a/src/awkward/index.py
+++ b/src/awkward/index.py
@@ -177,6 +177,32 @@ class Index:
def __len__(self) -> int:
return int(self.length)
+ _cycler = 0
+
+ def __getattr__(self, name):
+ if name == "__array_struct__":
+ self._cycler = 1
+ raise AttributeError(name)
+ elif name == "__array_interface__":
+ if self._cycler == 1:
+ self._cycler = 2
+ else:
+ self._cycler = 0
+ return self._data.__array_interface__
+ elif name == "__array__":
+ if self._cycler == 2:
+ self._cycler = 3
+ raise Exception(
+ "Module is coming to the conclusion that Index is a Python iterable"
+ )
+ else:
+ self._cycler = 0
+ raise AttributeError(name)
+ elif name in self.__dict__:
+ return self.__dict__[name]
+ else:
+ raise AttributeError(name)
+
@property
def __cuda_array_interface__(self):
return self._data.__cuda_array_interface__ # type: ignore[attr-defined]And if I run all the tests, we get this: So it looks like we iterate over virtual arrays pretty often by mistake. We should fix this with this PR ASAP. There are no failures with this PR. cc @ianna |
pfackeldey
left a comment
pfackeldey
left a comment
There was a problem hiding this comment.
Hah! Good catch @ikrommyd and thanks for the checks with the test suite.
Yes, that is important to get in!
It looks good to me, I'm not familiar with the cupy protocols, but the numpy and jax one look correct to me 👍 (I haven't seen __array_interface__, what's that used for?)
|
Are there actually cases where we want |
pfackeldey
left a comment
pfackeldey
left a comment
There was a problem hiding this comment.
Sorry, I need to request changes (withdraw my approval for now)...
I thought about it a bit more and I think we shouldn't add any of those protocols to VirtualNDArray. Instead, we need to trace down the cases where we forgot to call maybe_materialize (e.g. by erroring on __iter__).
What do you think about this?
I don't think this needs to be manual. In the same way virtual array + 1 automatically just works, I'd want
Well it's a buffer so yeah, if you want to manually iterate over a buffer, who are we to stop it. I'm not saying one should be iterating over arrays but yeah I'd rather allow it to be possible. |
|
Now that you mention that, I actually think the placeholder should implement these and throw the "good" error message that we have too. |
ikrommyd
changed the title
ikrommyd
changed the title
|
FWIW, I think It's actually not good that we have to add |
pfackeldey
left a comment
pfackeldey
left a comment
There was a problem hiding this comment.
We talked this through offline in detail:
These protocols are good to add, better than manual maybe_materialize. In the long run we can get rid of many of these maybe_materialize calls as well, which helps in maintaining this code. So all in all, I'm happy with these changes. Also, my concern of potential over materialization is not valid: we would've hit those over materializations already now through __iter__.
So, everything good from my side 👍
|
Thanks! Yeah I'll remove as many |
2 participants
This is to fix all the
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()issues we see with the cuda backend with__setitem__and__getitem__between virtual arrays with a cupy nplike and regular cupy arrays. We xfail several tests due to that. This will work with cupy 14 once it comes out and causes no errors otherwise. With cupy 14, the tests we xfail actually pass (tried with cupy installation from v14 pre-release) Once cupy 14 is out, we should remove those xfails too. See cupy/cupy#9089 and cupy/cupy#8352 and #3142 for more info.Also this is to fix a few more cases similar to the cupy issue but with the numpy backend where if you do
y[:] = xandyis a numpy array andxa virtual array, this will use__iter__and iterate over the virtual array like a python object. Adding the interface fixes that.