BUG: Setting an object array via slicing with nested lists invokes the objects `__array__` methods unnecessarily

[chriseclectic]

Describe the issue:

I have been running into issues with trying to use Numpy object arrays to contruct ND-Arrays of Python objects where I want to take advantage of the ND-indexing and other array features, but the values of the array are Python objects, which themselves have an __array__ method defined, but the __array__ method is inefficient and should be avoided.

In these cases I am not able to construct an object array from a nested list via slicing without invoking the objects __array__ method.

A simplified example of where this is a problem is as follows:

Consider a skeleton implementation of a Python class for N-qubit Pauli matrices. These objects have an efficient representation (for simplicity I will just use the length N string labels), but can also be represented inefficiently as shape (2 ** N, 2 **N) complex matrices. Suppose this class defines an __array__ method to convert to these inefficient complex arrays like shown in the bellow code example.

Now suppose I want to construct an object ND-array of Paulis. The simplest case would be

arr1 = np.empty([1], dtype=object)
arr1[0] = Pauli("X")

This does not invoke the Pauli.__array__ method and returns array([Pauli(X)], dtype=object).

However if I try and set via slicing

arr2 = np.empty([1], dtype=object)
arr2[:] = [Pauli("X")]
# Prints: Pauli.__array__ called

This invokes the array method, even though the returned array is an object array with the original objects as values (np.all(arr1 == arr2) is True).

Complete code shown bellow, where you can see where this becomes a problem for example by setting nq=20 which would require ~17TB memory per Pauli for each __array__ return.

Reproduce the code example:

import numpy as np

class Pauli:
    """Basic N-qubit Pauli"""

    _mats = {
        "I": np.eye(2, dtype=complex),
        "X": np.array([[0, 1], [1, 0]], dtype=complex),
        "Y": np.array([[0, -1j], [0, 1j]], dtype=complex),
        "Z": np.array([[1, 0], [0, -1]], dtype=complex),
    }

    def __init__(self, label):
        """Initialize a Pauli from a string label"""
        self.label = label
     
    def __repr__(self):
        return f"Pauli({self.label})"
    
    def to_matrix(self):
        """Convert to a (2**N, 2**N) complex matrix"""
        ret = np.eye(1, dtype=complex)
        for i in self.label:
            ret = np.kron(self._mats[i], ret)
        return ret

    def __array__(self, dtype=None):
        # Print to know when the `__array__` method is invoked
        print("Pauli.__array__ called")
        mat = self.to_matrix()
        if dtype:
            mat = mat.astype(dtype)
        return mat

    def __eq__(self, other):
        if isinstance(other, Pauli):
            return self.label == other.label
        return False

# Create an array-like nested list of Pauli objects
nq = 1
paulis = [[Pauli(nq * "I"), Pauli(nq * "X")],
          [Pauli(nq * "Y"), Pauli(nq * "Z")]]

# Initialize and empty array and set values to nested list elements
arr = np.empty([2, 2], dtype=object)
arr[:] = paulis

# This print `Pauli.__array__ called` 4 times, so invokes it for each Pauli
# Still returns an object array with the `Pauli` objects as values.

Error message:

No response

Runtime information:

My main numpy runtime is

1.23.5
3.9.13 | packaged by conda-forge | (main, May 27 2022, 17:01:00) 
[Clang 13.0.1 ]

I also tried on Numpy 1.25.2 and get same results.

Context for the issue:

This is probably quite a specific edge case, but it is something I've been struggling with for awhile and have had to come up with various messy work arounds to avoid the issue.

[rkern]

Assuming that you only have nested lists and your objects don't subclass from list, here is a somewhat robust workaround for the time being:

def flatten_lol(lol):
    for sub in lol:
        if isinstance(sub, list):
            yield from flatten_lol(sub)
        else:
            yield sub

def shape_lol(lol):
    child = lol[0]
    if isinstance(child, list):
        return (len(lol),) + shape_lol(child)
    else:
        return (len(lol),)

def array_from_lol(lol):
    shape = shape_lol(lol)
    arr = np.fromiter(flatten_lol(lol), dtype=object).reshape(shape)
    return arr

[seberg]

Agree that we should try to avoid this. One detail to think about is what happens in these two cases:

np.array([None, Pauli(...)], dtype=object)
np.array([Pauli(...), None], dtype=object)

are we OK with not calling __array__() in the first, but doing so in the second or should we be precise and only not call things when the user requested a maximum number of dimensions (which is the case here implicitly).

We have all the information needed to adjust this either way (or can add it easily). The code that does this is part of PyArray_DiscoverDTypeAndShape_Recursive in array_coercion.c.

[rkern]

To be clear, it's not np.array(nested_list) that's being reported here; that's probably hopeless. Rather arr[:] = nested_list, which is what we usually recommend as the workaround.

[seberg]

Yes, I am aware np.array(nested_list) doesn't have the maximum depth information while arr[...] = nested_list does. But, it does use the same machinery, effectively somthing like np.asarray(nested_list, dtype=arr.dtype, maxdims=arr.ndim).

[rkern]

In which case, no, I would expect no difference between the two cases.