search

scientific-python / spatch

BSD 3-Clause "New" or "Revised" License
2 stars 3 forks source link

Start of a prototype for string-based, very strict, type dispatching #2

Open seberg opened 2 months ago

seberg commented 2 months ago

This is an early (and very dirty) start, in case it is useful to get going. (Ping @eriknw)

The "example" is just the prototype_* modules, with the _example one being executable. I did not look into the entry-points, and of course this is all very raw:

  1. We can make the way we dispatch on types much smarter:
    • For speed
    • For backend order (potentially), this code doesn't even try to address that.
    • I had discussed with Eric, having types we expect/need to match, and types we understand can make sense. E.g. cucim may want to only kick in for cupy arrays, but understand numpy arrays. A better example maybe a dask.array backend which can understand cupy/numpy, but will never be picked for those.
  2. This is just string based, that can be relaxed of course.
  3. I added a WillNotHandle() return class to pass out information for rejection, since I liked that idea:
    • (We cannot "force" right now, in this way, I suspect that is OK for now.)
  4. I solved the parameter selection with strings + inspect, which I really like for simplicity. We need the signature fetching for it, but otherwise we can always optimize it.

Anyway, sharing in the hoping it might help. But please feel free to ignore if you have another start and are pushing forward there!

lucascolley commented 2 months ago

For comparison, this is how we're currently doing basic type-delegation in scipy.ndimage: https://github.com/scipy/scipy/blob/main/scipy/ndimage/_support_alternative_backends.py, https://github.com/scipy/scipy/blob/main/scipy/ndimage/_dispatchers.py.

seberg commented 2 months ago

FWIW, this is a branch in skimage: https://github.com/scikit-image/scikit-image/compare/main...seberg:backend-cucim-try that uses the state here. I added a backend for cucim there (for filters.median and filters.gaussian), that has to live in cucim and be auto-generated of course. Difference in speed for a 1000x1000 NumPy vs. CuPy array is, 114us vs. 31ms on the machine I tried.