Forwarding arguments to the right nested functions in python

Question

The dask apply_along_axis function has the following signature:

dask.array.apply_along_axis(func1d, axis, arr, *args, dtype=None, shape=None, **kwargs)

In my case, func1d has the following signature:

my_fun(arr, x, xp, propagate=True)

I need to specify dtype and shape arguments (their default to None does not fit my case). I'm not familiar with *args and **kwargs syntaxes (I come from C++ with template parameter packs syntaxes), and even if I've tried a number call forms, I always get something along my_fun() got an unexpected keyword argument 'shape'.

How do I call this function to pass the right arguments to the right functions?

Repoducible example:

import dask
import dask.array as da
import numpy as np

def my_fun(data, x, xp):
    return data

new_array = np.zeros((100,100,100))
big_array=da.from_array(new_array, chunks=(100,100,100))
x="foo"
xp="fee"
interpolated = da.apply_along_axis(func1d=my_fun, axis=0, arr=big_array, shape=big_array.shape, dtype=big_array.dtype, x=x, xp=xp).compute()

Returns:

Traceback (most recent call last):
  File "parallelDask.py", line 28, in <module>
    interpolated = da.apply_along_axis(func1d=my_fun, axis=0, arr=big_array, shape=big_array.shape, dtype=big_array.dtype, x=x, xp=xp).compute()
  File "/home/becheler/dev/virtual_environments/crumbs-env/lib/python3.8/site-packages/dask/array/routines.py", line 304, in apply_along_axis
    test_result = np.array(func1d(test_data, *args, **kwargs))
TypeError: my_fun() got an unexpected keyword argument 'shape'

EDIT: I think it comes pretty close to this issue but I could not figure out how to solve the problem.

Answer 1

Updating the Dask version fixes this issue. We've verified this works with Dask 2022.3.0 -- latest Dask version at the time of answering. :)