Dynamic Inheritance for CPU/GPU-Aware Classes in Python

Dynamic Array Handling with Conditional Inheritance in Python

This solution demonstrates dynamic inheritance in Python using NumPy and CuPy for CPU/GPU-agnostic array handling. It employs Python's object-oriented programming for flexibility and modularity.

from typing import Union
import numpy as np
import cupy as cp
# Base class for shared functionality
class BaseArray:
    def bar(self, x):
        # Example method: Add x to the array
        return self + x
# Numpy-specific class
class NumpyArray(BaseArray, np.ndarray):
    pass
# CuPy-specific class
class CuPyArray(BaseArray, cp.ndarray):
    pass
# Factory function to handle conditional inheritance
def create_array(foo: Union[np.ndarray, cp.ndarray]):
    if isinstance(foo, cp.ndarray):
        return foo.view(CuPyArray)
    return foo.view(NumpyArray)
# Example usage
if __name__ == "__main__":
    foo_np = np.array([1.0, 2.0, 3.0])
    foo_cp = cp.array([1.0, 2.0, 3.0])
    array_np = create_array(foo_np)
    array_cp = create_array(foo_cp)
    print(array_np.bar(2))  # [3.0, 4.0, 5.0]
    print(array_cp.bar(2))  # [3.0, 4.0, 5.0] (on GPU)

Alternative Approach Using Class Wrapping

This solution uses a wrapper class to dynamically delegate CPU/GPU behavior based on the input type. The focus is on clean code and separation of concerns.

from typing import Union
import numpy as np
import cupy as cp
# Wrapper class for CPU/GPU agnostic operations
class ArrayWrapper:
    def __init__(self, foo: Union[np.ndarray, cp.ndarray]):
        self.xp = cp.get_array_module(foo)
        self.array = foo
    def add(self, value):
        return self.xp.array(self.array + value)
# Example usage
if __name__ == "__main__":
    foo_np = np.array([1.0, 2.0, 3.0])
    foo_cp = cp.array([1.0, 2.0, 3.0])
    wrapper_np = ArrayWrapper(foo_np)
    wrapper_cp = ArrayWrapper(foo_cp)
    print(wrapper_np.add(2))  # [3.0, 4.0, 5.0]
    print(wrapper_cp.add(2))  # [3.0, 4.0, 5.0] (on GPU)

Unit Tests for Both Solutions

Unit tests to ensure the solutions work as expected across CPU and GPU environments.

import unittest
import numpy as np
import cupy as cp
class TestArrayInheritance(unittest.TestCase):
    def test_numpy_array(self):
        foo = np.array([1.0, 2.0, 3.0])
        array = create_array(foo)
        self.assertTrue(isinstance(array, NumpyArray))
        self.assertTrue(np.array_equal(array.bar(2), np.array([3.0, 4.0, 5.0])))
    def test_cupy_array(self):
        foo = cp.array([1.0, 2.0, 3.0])
        array = create_array(foo)
        self.assertTrue(isinstance(array, CuPyArray))
        self.assertTrue(cp.array_equal(array.bar(2), cp.array([3.0, 4.0, 5.0])))
if __name__ == "__main__":
    unittest.main()

Essential Questions About Dynamic Inheritance in Python

1. What is dynamic inheritance?
2. Dynamic inheritance allows a class to adjust its behavior or parent class at runtime based on input, like switching between NumPy and CuPy.
3. How does get_array_module work?
4. This CuPy function determines whether an array is a NumPy or CuPy instance, enabling backend selection for operations.
5. What is the role of view() in inheritance?
6. The view() method in both NumPy and CuPy creates a new array instance with the same data but assigns it a different class.
7. How does dynamic inheritance improve performance?
8. By selecting optimized backends and avoiding redundant logic, dynamic inheritance ensures efficient CPU and GPU utilization.
9. Can I add additional backends in the future?
10. Yes, by designing your dynamic inheritance logic modularly, you can include libraries like TensorFlow or JAX without rewriting existing code.

Key Takeaways for Effective Dynamic Inheritance

Dynamic inheritance in Python provides a powerful way to create flexible and hardware-agnostic classes. By choosing modular and efficient designs, you ensure that your code remains maintainable while adapting to different backends like NumPy and CuPy. This versatility benefits projects requiring scalability and performance.

Incorporating solutions like those demonstrated in this article allows developers to focus on solving domain-specific challenges. Real-world examples, such as transitioning from CPU prototypes to GPU-heavy workloads, highlight the importance of adaptable code. With these principles, dynamic inheritance becomes a cornerstone of robust Python programming. 💡