# Buggy code
def _apply_with_dask_atop(
    func, args, input_dims, output_dims, signature, output_dtypes, output_sizes=None
):
    import dask.array as da

    if signature.num_outputs > 1:
        raise NotImplementedError(
            "multiple outputs from apply_ufunc not yet "
            "supported with dask='parallelized'"
        )

    if output_dtypes is None:
        raise ValueError(
            "output dtypes (output_dtypes) must be supplied to "
            "apply_func when using dask='parallelized'"
        )
    if not isinstance(output_dtypes, list):
        raise TypeError(
            "output_dtypes must be a list of objects coercible to "
            "numpy dtypes, got {}".format(output_dtypes)
        )
    if len(output_dtypes) != signature.num_outputs:
        raise ValueError(
            "apply_ufunc arguments output_dtypes and "
            "output_core_dims must have the same length: {} vs {}".format(
                len(output_dtypes), signature.num_outputs
            )
        )
    (dtype,) = output_dtypes

    if output_sizes is None:
        output_sizes = {}

    new_dims = signature.all_output_core_dims - signature.all_input_core_dims
    if any(dim not in output_sizes for dim in new_dims):
        raise ValueError(
            "when using dask='parallelized' with apply_ufunc, "
            "output core dimensions not found on inputs must have "
            "explicitly set sizes with ``output_sizes``: {}".format(new_dims)
        )

    for n, (data, core_dims) in enumerate(zip(args, signature.input_core_dims)):
        if isinstance(data, dask_array_type):
            # core dimensions cannot span multiple chunks
            for axis, dim in enumerate(core_dims, start=-len(core_dims)):
                if len(data.chunks[axis]) != 1:
                    raise ValueError(
                        "dimension {!r} on {}th function argument to "
                        "apply_ufunc with dask='parallelized' consists of "
                        "multiple chunks, but is also a core dimension. To "
                        "fix, rechunk into a single dask array chunk along "
                        "this dimension, i.e., ``.rechunk({})``, but beware "
                        "that this may significantly increase memory usage.".format(
                            dim, n, {dim: -1}
                        )
                    )

    (out_ind,) = output_dims
    # skip leading dimensions that we did not insert with broadcast_compat_data
    atop_args = [
        element
        for (arg, dims) in zip(args, input_dims)
        for element in (arg, dims[-getattr(arg, "ndim", 0) :])
    ]
    return da.atop(
        func, out_ind, *atop_args, dtype=dtype, concatenate=True, new_axes=output_sizes
    )


# Traceback (IndexError)
xarray.apply_ufunc(
scipy.stats.norm.cdf, a.chunk(), 0, 1, dask='parallelized', output_dtypes=[a.dtype]
).compute()

IndexError                                Traceback (most recent call last)
<ipython-input-35-1d4025e1ebdb> in <module>()
----> 1 xarray.apply_ufunc(scipy.stats.norm.cdf, a.chunk(), 0, 1, dask='parallelized', output_dtypes=[a.dtype]).compute()

~/anaconda3/lib/python3.6/site-packages/xarray/core/computation.py in apply_ufunc(func, *args, **kwargs)
913                                      join=join,
914                                      exclude_dims=exclude_dims,
--> 915                                      keep_attrs=keep_attrs)
916     elif any(isinstance(a, Variable) for a in args):
917         return variables_ufunc(*args)

~/anaconda3/lib/python3.6/site-packages/xarray/core/computation.py in apply_dataarray_ufunc(func, *args, **kwargs)
210
211     data_vars = [getattr(a, 'variable', a) for a in args]
--> 212     result_var = func(*data_vars)
213
214     if signature.num_outputs > 1:

~/anaconda3/lib/python3.6/site-packages/xarray/core/computation.py in apply_variable_ufunc(func, *args, **kwargs)
561             raise ValueError('unknown setting for dask array handling in '
562                              'apply_ufunc: {}'.format(dask))
--> 563     result_data = func(*input_data)
564
565     if signature.num_outputs > 1:

~/anaconda3/lib/python3.6/site-packages/xarray/core/computation.py in <lambda>(*arrays)
555             func = lambda *arrays: _apply_with_dask_atop(
556                 numpy_func, arrays, input_dims, output_dims, signature,
--> 557                 output_dtypes, output_sizes)
558         elif dask == 'allowed':
559             pass

~/anaconda3/lib/python3.6/site-packages/xarray/core/computation.py in _apply_with_dask_atop(func, args, input_dims, output_dims, signature, output_dtypes, output_sizes)
624                  for element in (arg, dims[-getattr(arg, 'ndim', 0):])]
625     return da.atop(func, out_ind, *atop_args, dtype=dtype, concatenate=True,
--> 626                    new_axes=output_sizes)
627
628

~/anaconda3/lib/python3.6/site-packages/dask/array/core.py in atop(func, out_ind, *args, **kwargs)
2231         raise ValueError("Must specify dtype of output array")
2232
-> 2233     chunkss, arrays = unify_chunks(*args)
2234     for k, v in new_axes.items():
2235         chunkss[k] = (v,)

~/anaconda3/lib/python3.6/site-packages/dask/array/core.py in unify_chunks(*args, **kwargs)
2117             chunks = tuple(chunkss[j] if a.shape[n] > 1 else a.shape[n]
2118                            if not np.isnan(sum(chunkss[j])) else None
-> 2119                            for n, j in enumerate(i))
2120             if chunks != a.chunks and all(a.chunks):
2121                 arrays.append(a.rechunk(chunks))

~/anaconda3/lib/python3.6/site-packages/dask/array/core.py in <genexpr>(.0)
2117             chunks = tuple(chunkss[j] if a.shape[n] > 1 else a.shape[n]
2118                            if not np.isnan(sum(chunkss[j])) else None
-> 2119                            for n, j in enumerate(i))
2120             if chunks != a.chunks and all(a.chunks):
2121                 arrays.append(a.rechunk(chunks))

IndexError: tuple index out of range

# Fixed code
def _apply_with_dask_atop(
    func, args, input_dims, output_dims, signature, output_dtypes, output_sizes=None
):
    import dask.array as da

    if signature.num_outputs > 1:
        raise NotImplementedError(
            "multiple outputs from apply_ufunc not yet "
            "supported with dask='parallelized'"
        )

    if output_dtypes is None:
        raise ValueError(
            "output dtypes (output_dtypes) must be supplied to "
            "apply_func when using dask='parallelized'"
        )
    if not isinstance(output_dtypes, list):
        raise TypeError(
            "output_dtypes must be a list of objects coercible to "
            "numpy dtypes, got {}".format(output_dtypes)
        )
    if len(output_dtypes) != signature.num_outputs:
        raise ValueError(
            "apply_ufunc arguments output_dtypes and "
            "output_core_dims must have the same length: {} vs {}".format(
                len(output_dtypes), signature.num_outputs
            )
        )
    (dtype,) = output_dtypes

    if output_sizes is None:
        output_sizes = {}

    new_dims = signature.all_output_core_dims - signature.all_input_core_dims
    if any(dim not in output_sizes for dim in new_dims):
        raise ValueError(
            "when using dask='parallelized' with apply_ufunc, "
            "output core dimensions not found on inputs must have "
            "explicitly set sizes with ``output_sizes``: {}".format(new_dims)
        )

    for n, (data, core_dims) in enumerate(zip(args, signature.input_core_dims)):
        if isinstance(data, dask_array_type):
            # core dimensions cannot span multiple chunks
            for axis, dim in enumerate(core_dims, start=-len(core_dims)):
                if len(data.chunks[axis]) != 1:
                    raise ValueError(
                        "dimension {!r} on {}th function argument to "
                        "apply_ufunc with dask='parallelized' consists of "
                        "multiple chunks, but is also a core dimension. To "
                        "fix, rechunk into a single dask array chunk along "
                        "this dimension, i.e., ``.rechunk({})``, but beware "
                        "that this may significantly increase memory usage.".format(
                            dim, n, {dim: -1}
                        )
                    )

    (out_ind,) = output_dims

    atop_args = []
    for arg, dims in zip(args, input_dims):
        # skip leading dimensions that are implicitly added by broadcasting
        ndim = getattr(arg, "ndim", 0)
        trimmed_dims = dims[-ndim:] if ndim else ()
        atop_args.extend([arg, trimmed_dims])

    return da.atop(
        func, out_ind, *atop_args, dtype=dtype, concatenate=True, new_axes=output_sizes
    )