thunder-project/thunder

dtype propogation in images.map() operation

boazmohar opened this issue · 3 comments

I am trying to do a map opperation that changes my value size and I know the resulting size.
Following the code I see that in thunder images map() exposes dims but not dtype:

def map(self, func, dims=None, with_keys=False):

It callse _map from the base class:

return self._map(func, axis=0, value_shape=dims, with_keys=with_keys)

without passing dtype.

_map() does accept a dtype that defaults to None:

_map(self, func, axis=(0,), value_shape=None, dtype=None, with_keys=False):

calling the bolt map operation always with None:

mapped = self.values.map(expand, axis, value_shape, dtype, with_keys)

In Bolt, if either dims which is now value_shape or dtype are None it will test the operation with a random array:

if value_shape is None or dtype is None:
            # try to compute the size of each mapped element by applying func to a random array
            try:
                mapped = test_func(random.randn(*swapped.values.shape).astype(self.dtype))
            except Exception:
                first = swapped._rdd.first()
                if first:
                    # eval func on the first element
                    mapped = test_func(first[1])
            if value_shape is None:
                value_shape = mapped.shape
            if dtype is None:
                dtype = mapped.dtype

I will be happy to do a PR if anyone is willing to show me how (it will be my first)

Boaz

@boazmohar, great catch! Happy to show you how to submit a PR with your enhancement 👍

@jwittenbach Thanks, will you be at Bob's at around 16:30?