I'm +1 on this approach. In general I think that we should leave kwarg checking to blockwise when it tries out the function on a small sample.

I'm +1 on this approach. In general I think that we should leave kwarg checking to blockwise when it tries out the function on a small sample.

Me too, this will make things much simpler for __array_function__ going further.

Thanks for putting this together Peter. Looks great! 😄

It seems this causes test_einsum_invalid_args to fail.

Thanks @jakirkham for pointing it out, indeed I had not seen it failed. I'll work on it.

Just a quick update on this: the issue is a bit more complex than I first predicted, blockwise_meta is shadowing exceptions from function calls to other libraries (in this case, np.einsum), so I'm working on a fix for that.

Right, so we may or may not want to raise a TypeError on failed calls on meta. This is a complex question.

Fortunately we've already been doing this in Dask Dataframe for a year or two now, so we have some information on how what happens. In general we get lots of user bug reports on malformed meta. We try functions on tiny sample functions and they fail, and users get pretty frustrated. Most of them don't understand things well enough to understand how to fix their problem. In Dask Dataframe we kind of have to maintain meta in order to operate effectively (we need column names and dtypes more than we do in Dask Array), so we're sort of stuck with having to explain to people, over and over, that they need to provide proper metadata. This is, as you can imagine, super-frustrating for all involved.

In the Dask array case we might want to be a bit more lenient. It seems like there is a spectrum of strictness that we could use:

1. Not strict: default back to a dtype transform, or even just return the meta of the first input array
2. Warn: Raise a warning that things didn't work out great, inform them how to remove the warning, and probably point to a longer doc page on the topic. However continue on with the not strict approach and give them a possibly malformed dask array.
3. Raise: Same as above, but we raise and don't allow them to continue.

I actually realized one thing on the error here, the test currently is:

     with pytest.raises(TypeError):
         da.einsum('a', *da_inputs, foo=1, bar=2)

and Dask raises the exception. However, we're now deferring the exception to the compute library (NumPy, CuPy), so what we really need to do is to add .compute() at the end, rather than making blockwise_meta raise the exception. I guess that makes sense, what do you think @jakirkham, @mrocklin ?

In fact, here's one other existing test:

def test_average_raises():
    d_a = da.arange(11, chunks=2)

    with pytest.raises(TypeError):
        da.average(d_a, weights=[1, 2, 3])

    with pytest.warns(RuntimeWarning):
        da.average(d_a, weights=da.zeros_like(d_a)).compute()

in this case, Dask raises a TypeError, but NumPy raises a RuntimeWarning, which is what I think we should expect for the einsum case failing in this PR too.

I actually realized one thing on the error here, the test currently is:

     with pytest.raises(TypeError):
         da.einsum('a', *da_inputs, foo=1, bar=2)

and Dask raises the exception. However, we're now deferring the exception to the compute library (NumPy, CuPy), so what we really need to do is to add .compute() at the end, rather than making blockwise_meta raise the exception. I guess that makes sense, what do you think @jakirkham, @mrocklin ?

I think that it's still good to catch errors at graph construction time rather than compute time if we can. Obviously, this isn't always possible, but I think that we can do it in the case of bad keywords.

I recommend that we make a function like this, and use it whenever we try to get meta from calling a user provided function:

def compute_meta(func, *args, **kwargs):
    bad_keywords = set(inspect.signature(func).parameters) - set(kwargs)
    if bad_keywords:
        raise TypeError("Unexpected keywords to %s function: %s" % (funcname(func), bad_keywords)

    args = [getattr(arg, "_meta", arg) for arg in args]
    kwargs = {k: getattr(v, "_meta", v) for k, v in kwargs.items()}
    
    try:
        meta = func(*args, **kwargs)
    except Exception:
        meta = ...

    return meta

@mrocklin we can definitely do that too for this particular case. But this will clearly not be the case for all exceptions, and I think the more we move towards __array_function__, it's likely that we're going to push more of such errors to the compute task itself. Are we ok with that?

@mrocklin we can definitely do that too for this particular case. But this will clearly not be the case for all exceptions, and I think the more we move towards array_function, it's likely that we're going to push more of such errors to the compute task itself. Are we ok with that?

To a certain extent I think that it's necessary. People will provide functions that strongly expect certain things that won't be true of meta. I recommend that we also include the following test.

def test_permissible_meta_calculations():
    def f(x):
        assert x.shape == (10, 10)
        return x + 1

    x = da.ones((30, 30), chunks=(10, 10))
    y = x.map_blocks(f)

    assert_eq(y, np.ones((30, 30)) + 1)

Once again, this is a nice solution, except it has a lot of caveats. I've attempted to work on your suggestion @mrocklin, but it's not so simple. Again we fall on the shadowing of exceptions that I've mentioned before, but another issue comes from functions that take **kwargs instead of the actual named parameters, where we can't test for a bad signature.

I think we should then move conservatively (from the _meta perspective), just like we are doing in #4938, and for now let this fail during compute. Do you have any objection here?

I will raise a few issues (and perhaps some [WIP] PRs) with a few ideas I have attempted that fix some issues, but not all of them, and then we can build upon those ideas.

Yes, I agree that if the function has **kwargs then the keyword check should just pass.

I think it's suboptimal to fail on compute but probably necessary some of the time. I'm curious what we did before.

In [1]: import dask; dask.__version__
Out[1]: '1.2.2'

In [2]: import dask.array as da

In [3]: x = da.ones((30, 30), chunks=(10, 10))

In [4]: def f(x):
   ...:     assert x.shape == (10, 10)
   ...:     return x + 1
   ...:

In [5]: x.map_blocks(f)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-5f8635e7db98> in <module>
----> 1 x.map_blocks(f)

~/workspace/dask/dask/array/core.py in map_blocks(self, func, *args, **kwargs)
   1637     @wraps(map_blocks)
   1638     def map_blocks(self, func, *args, **kwargs):
-> 1639         return map_blocks(func, self, *args, **kwargs)
   1640
   1641     def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):

~/workspace/dask/dask/array/core.py in map_blocks(func, *args, **kwargs)
    493
    494     if dtype is None:
--> 495         dtype = apply_infer_dtype(func, args, original_kwargs, 'map_blocks')
    496
    497     if drop_axis:

~/workspace/dask/dask/array/core.py in apply_infer_dtype(func, args, kwargs, funcname, suggest_dtype, nout)
    282         msg = None
    283     if msg is not None:
--> 284         raise ValueError(msg)
    285     return o.dtype if nout is None else tuple(e.dtype for e in o)
    286

ValueError: `dtype` inference failed in `map_blocks`.

Please specify the dtype explicitly using the `dtype` kwarg.

Original error is below:
------------------------
AssertionError()

Traceback:
---------
  File "/Users/mrocklin/workspace/dask/dask/array/core.py", line 267, in apply_infer_dtype
    o = func(*args, **kwargs)
  File "<ipython-input-4-014b0f39d8bc>", line 2, in f
    assert x.shape == (10, 10)

Hrm, so it looks like what I'm proposing is a relaxation of our previous policy, where we were more strict. It still feels correct to me to loosen things here, mostly because I don't trust users to understand and provide a meta array as much as I trust them to understand and provide a dtype.

Lets cc @jcrist and @shoyer to get their thoughts here. What should we do in the following situation where our attempt to run a user defined function fails on a _meta chunk. Should we err, warn, pass? Discussion above.

def test_permissible_meta_calculations():
    def f(x):
        assert x.shape == (10, 10)
        return x + 1

    x = da.ones((30, 30), chunks=(10, 10))
    y = x.map_blocks(f)

    assert_eq(y, np.ones((30, 30)) + 1)

Right now, blockwise_meta ignores a lot of TypeErrors (by catching Exception) and defaults to constructing dask.array with the dtype, so that's the major problem when we let TypeError pass. I do think we should increase coverage of possible cases here, but I also think that will take time and we'll break more user code in the short term, so that's why I'm proposing a short-term solution for this PR and continue the discussion on how we increase the _meta coverage for this situation.

If you'd like a strict mode so that it's easier for developers to detect issues we would always add a strict mode relatively easily and control it with configuration

Yes, I agree with that, mainly because I don't think we will fix appropriately all issues we get on the tests in a timely manner. I will raise various issues tomorrow for that and the things I mentioned before.

So for now my question is: are you ok if we just call .compute() for the failing test here as a short-term solution and we continue with a more detailed conversation on the new issues?

So for now my question is: are you ok if we just call .compute() for the failing test here as a short-term solution and we continue with a more detailed conversation on the new issues?

For this test in particular I think that we should not call compute. I think that we should implement the inspect solution. That's something I think we can do in any case. I'm happy to push a commit like this if you prefer.

What do you mean with the "inspect solution"?

Ah, nevermind, I understand now. But this is what I'm saying: we have to allow TypeErrors to pass, and with that we have lots of other errors, not just the kwargs error.

I think that short term we should keep the except TypeError in the code. I think that we should check for keyword arguments, which will make the current failing test pass (and improve user experience I think). I think that this would be a good incremental change while we figure out what to do next. If there are other failing tests that would result from this approach I'm not aware of them (there appears to be only one at the moment, and I think we can fix it quickly), but would not be surprised if there are other things going on with cupy or whatnot that I'm not aware of.

(also, I'm shutting down for the night, so may not repond until much later)

That's what I'm saying: we fix that one test, yes, but dozens of others fail, not because of argument, but because we now let TypeError pass and they raise TypeErrors for various different reasons, and this is what except Exception: is preventing when it falls back to dtype.

Long story short: we fix this one issue and raise dozens of others. So I stand by my proposal: use compute for now until we can solve the dozens of other errors, which will require some time and perhaps a more consistent approach to operating on the 0-length meta arrays in blockwise_meta.

Do you mind if I push a small commit to your branch with my proposal? I think that this will be faster than prose communication.

Do it.

Ah, I see now. We're using a custom chunk_einsum function that has **kwargs, so my suggestion of using the inspect module doesn't suffice to detect the failure early.

Yup. Sorry, this was my mistake, we've talked offline on how we could fix this case for einsum, but I should have pointed it out here as well.

Yes, so I now agree with you that, given that we're using **kwargs in einsum, we should probably call .compute in the test. This is a bit of a regression in usability, but it seems like the right choice to me at the moment.

@mrocklin we can definitely undo that in the future, but we have to improve _meta exception handling before.

Thanks @pentschev ! Sorry for getting in the way here for a while :)

I've created #4951 and #4952 summarizing the discussion we had here, that would hopefully help us improve _meta.

Would it make sense to still raise TypeError if unrecognized keyword arguments that NumPy doesn't know how to handle are used?

e.g.,

if not kwargs.keys() <= {'casting', order'}:
    raise TypeError

We still don't have to pass them on to np.einsum() unless they were explicitly provided.

Xarray has a unit-test that verifies that the errors get raised for invalid keyword arguments (pydata/xarray#3009). We could just remove the test, but it seems like a reasonable thing to keep around.

@shoyer https://github.com/dask/dask/pull/4954/files#diff-3598e11c3b3a16881fecff61abdf1f4cR97 will solve that.

It's worth noting that np.einsum itself doesn't note which arguments it takes in its signature, but only takes a generic **kwargs. This makes it hard for us to inspect and check validity of arguments at this level.

In [1]: import numpy as np                                                      

In [2]: import inspect                                                          

In [3]: inspect.getfullargspec(np.einsum)                                       
Out[3]: FullArgSpec(args=[], varargs='operands', varkw='kwargs', defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={})

It would be better if NumPy's einsum explicitly listed keyword arguments it would accept so we could inspect it and stay current. Explicit checking on our end will inevitably go stale if the allowed keyword arguments change in the future.

As @pentschev mentions above, this is being resolved in #4954

Please excuse our churn here.

@shoyer so far I think that we'll be able to make the next Dask release fully backwards compatible from an Xarray perspective.