Subclassing Numpy Array - Propagate Attributes

import numpy as np

class ExampleTensor(np.ndarray):
    def __new__(cls, input_array):
        return np.asarray(input_array).view(cls)

    def __array_finalize__(self, obj) -> None:
        if obj is None: return
        # This attribute should be maintained!
        default_attributes = {"attr": 1}
        self.__dict__.update(default_attributes)  # another way to set attributes

Implement the array_ufunc method like this

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):  # this method is called whenever you use a ufunc
        f = {
            "reduce": ufunc.reduce,
            "accumulate": ufunc.accumulate,
            "reduceat": ufunc.reduceat,
            "outer": ufunc.outer,
            "at": ufunc.at,
            "__call__": ufunc,
        }
        output = ExampleTensor(f[method](*(i.view(np.ndarray) for i in inputs), **kwargs))  # convert the inputs to np.ndarray to prevent recursion, call the function, then cast it back as ExampleTensor
        output.__dict__ = self.__dict__  # carry forward attributes
        return output

Test

x = ExampleTensor(np.array([1,2,3]))
x.attr = 2

y0 = np.add(x, x)
print(y0, y0.attr)
y1 = np.add.outer(x, x)
print(y1, y1.attr)  # works even if called with method

[2 4 6] 2
[[2 3 4]
 [3 4 5]
 [4 5 6]] 2

Explanation in the comments.


Here is an attempt that works for operators that are not arrays and even when our subclass is specified as output of a numpy ufunc (explanations in the comments):

import numpy as np


class ArraySubclass(np.ndarray):
    '''Subclass of ndarray MUST be initialized with a numpy array as first argument.
    '''
    def __new__(cls, input_array, a=None, b=1):
        obj = np.asarray(input_array).view(cls)
        obj.a = a
        obj.b = b
        return obj

    def __array_finalize__(self, obj):
        if obj is None:  # __new__ handles instantiation
            return
        '''we essentially need to set all our attributes that are set in __new__ here again (including their default values). 
        Otherwise numpy's view-casting and new-from-template mechanisms would break our class.
        '''
        self.a = getattr(obj, 'a', None)
        self.b = getattr(obj, 'b', 1)

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):  # this method is called whenever you use a ufunc
        '''this implementation of __array_ufunc__ makes sure that all custom attributes are maintained when a ufunc operation is performed on our class.'''

        # convert inputs and outputs of class ArraySubclass to np.ndarray to prevent infinite recursion
        args = ((i.view(np.ndarray) if isinstance(i, ArraySubclass) else i) for i in inputs)
        outputs = kwargs.pop('out', None)
        if outputs:
            kwargs['out'] = tuple((o.view(np.ndarray) if isinstance(o, ArraySubclass) else o) for o in outputs)
        else:
            outputs = (None,) * ufunc.nout
        # call numpys implementation of __array_ufunc__
        results = super().__array_ufunc__(ufunc, method, *args, **kwargs)  # pylint: disable=no-member
        if results is NotImplemented:
            return NotImplemented
        if method == 'at':
            # method == 'at' means that the operation is performed in-place. Therefore, we are done.
            return
        # now we need to make sure that outputs that where specified with the 'out' argument are handled corectly:
        if ufunc.nout == 1:
            results = (results,)
        results = tuple((self._copy_attrs_to(result) if output is None else output)
                        for result, output in zip(results, outputs))
        return results[0] if len(results) == 1 else results

    def _copy_attrs_to(self, target):
        '''copies all attributes of self to the target object. target must be a (subclass of) ndarray'''
        target = target.view(ArraySubclass)
        try:
            target.__dict__.update(self.__dict__)
        except AttributeError:
            pass
        return target

and here are the corresponding unittests:

import unittest
class TestArraySubclass(unittest.TestCase):
    def setUp(self):
        self.shape = (10, 2, 5)
        self.subclass = ArraySubclass(np.zeros(self.shape))

    def test_instantiation(self):
        self.assertIsInstance(self.subclass, np.ndarray)
        self.assertIs(self.subclass.a, None)
        self.assertEqual(self.subclass.b, 1)
        self.assertEqual(self.subclass.shape, self.shape)
        self.assertTrue(np.array_equal(self.subclass, np.zeros(self.shape)))
        sub2 = micdata.arrayasubclass.ArraySubclass(np.zeros(self.shape), a=2)
        self.assertEqual(sub2.a, 2)

    def test_view_casting(self):
        self.assertIsInstance(np.zeros(self.shape).view(ArraySubclass),ArraySubclass)

    def test_new_from_template(self):
        self.subclass.a = 5
        bla = self.subclass[3, :]
        self.assertIsInstance(bla, ArraySubclass)
        self.assertIs(bla.a, 5)
        self.assertEqual(bla.b, 1)

    def test_np_min(self):
        self.assertEqual(np.min(self.subclass), 0)

    def test_ufuncs(self):
        self.subclass.b = 2
        self.subclass += 2
        self.assertTrue(np.all(self.subclass == 2))
        self.subclass = self.subclass + np.ones(self.shape)
        self.assertTrue(np.all(self.subclass == 3))
        np.multiply.at(self.subclass, slice(0, 2), 2)
        self.assertTrue(np.all(self.subclass[:2] == 6))
        self.assertTrue(np.all(self.subclass[2:] == 3))
        self.assertEqual(self.subclass.b, 2)

    def test_output(self):
        self.subclass.a = 3
        bla = np.ones(self.shape)
        bla *= 2
        np.multiply(bla, bla, out=self.subclass)
        self.assertTrue(np.all(self.subclass == 5))
        self.assertEqual(self.subclass.a, 3)

P.s. tempname123 got it almost right. However, his answer fails for operators that are not arrays and when his class is specified as output of a ufunc:

>>> ExampleTensor += 1
AttributeError: 'int' object has no attribute 'view'
>>> np.multiply(np.ones((5)), np.ones((5)), out=ExampleTensor)
RecursionError: maximum recursion depth exceeded in comparison