# Buggy code
def apply_variable_ufunc(func, *args, **kwargs):
    """apply_variable_ufunc(func, *args, signature, exclude_dims=frozenset())"""
    from .variable import Variable

    signature = kwargs.pop("signature")
    exclude_dims = kwargs.pop("exclude_dims", _DEFAULT_FROZEN_SET)
    dask = kwargs.pop("dask", "forbidden")
    output_dtypes = kwargs.pop("output_dtypes", None)
    output_sizes = kwargs.pop("output_sizes", None)
    keep_attrs = kwargs.pop("keep_attrs", False)
    if kwargs:
        raise TypeError(
            "apply_variable_ufunc() got unexpected keyword arguments: %s" % list(kwargs)
        )

    dim_sizes = unified_dim_sizes(
        (a for a in args if hasattr(a, "dims")), exclude_dims=exclude_dims
    )
    broadcast_dims = tuple(
        dim for dim in dim_sizes if dim not in signature.all_core_dims
    )
    output_dims = [broadcast_dims + out for out in signature.output_core_dims]

    input_data = [
        broadcast_compat_data(arg, broadcast_dims, core_dims)
        if isinstance(arg, Variable)
        else arg
        for arg, core_dims in zip(args, signature.input_core_dims)
    ]

    if any(isinstance(array, dask_array_type) for array in input_data):
        if dask == "forbidden":
            raise ValueError(
                "apply_ufunc encountered a dask array on an "
                "argument, but handling for dask arrays has not "
                "been enabled. Either set the ``dask`` argument "
                "or load your data into memory first with "
                "``.load()`` or ``.compute()``"
            )
        elif dask == "parallelized":
            input_dims = [broadcast_dims + dims for dims in signature.input_core_dims]
            numpy_func = func

            def func(*arrays):
                return _apply_with_dask_atop(
                    numpy_func,
                    arrays,
                    input_dims,
                    output_dims,
                    signature,
                    output_dtypes,
                    output_sizes,
                )
        elif dask == "allowed":
            pass
        else:
            raise ValueError(
                "unknown setting for dask array handling in apply_ufunc: {}".format(
                    dask
                )
            )
    result_data = func(*input_data)

    if signature.num_outputs > 1:
        output = []
        for dims, data in zip(output_dims, result_data):
            var = Variable(dims, data)
            if keep_attrs and isinstance(args[0], Variable):
                var.attrs.update(args[0].attrs)
            output.append(var)
        return tuple(output)
    else:
        (dims,) = output_dims
        var = Variable(dims, result_data)
        if keep_attrs and isinstance(args[0], Variable):
            var.attrs.update(args[0].attrs)
        return var


# Traceback (ValueError)
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# function to apply\n",
    "A = np.random.rand(6, 4)\n",
    "apply_A = lambda data: A.dot(data.T).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 4)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# input data\n",
    "data = np.random.rand(2, 4)\n",
    "data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 6)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# change the core dimension from 4 to 6\n",
    "apply_A(data).shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (time: 2, x: 4)>\n",
       "array([[ 0.840646,  0.330686,  0.669603,  0.840243],\n",
       "       [ 0.419193,  0.936886,  0.189482,  0.543247]])\n",
       "Coordinates:\n",
       "  * time     (time) int64 0 1\n",
       "  * x        (x) int64 0 1 2 3"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# input DataArray\n",
    "dr = xr.DataArray(data, \n",
    "                  dims=['time', 'x'],\n",
    "                  coords=[np.arange(2), np.arange(4)],\n",
    "                  name='data'\n",
    "                 )\n",
    "dr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 1: not specifying input/output dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (time: 2, x: 6)>\n",
       "array([[ 1.24765 ,  0.24795 ,  1.053913,  1.338334,  1.9124  ,  0.860453],\n",
       "       [ 0.923691,  0.432066,  0.751094,  1.274393,  1.106883,  0.899529]])\n",
       "Coordinates:\n",
       "  * time     (time) int64 0 1\n",
       "  * x        (x) int64 0 1 2 3"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# note that the coordinate 'x' has incorrect dimension (should be 6, not 4)\n",
    "dr_out = xr.apply_ufunc(apply_A, dr)\n",
    "dr_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "conflicting sizes for dimension 'x': length 6 on 'data' and length 4 on 'x'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-4aafe6266f8a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# inconsistent coordinate size is not allowed in a NetCDF file\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mdr_out\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_netcdf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'apply_ufunc_result.nc'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/Research/Computing/miniconda3/envs/geo/lib/python3.6/site-packages/xarray/core/dataarray.py\u001b[0m in \u001b[0;36mto_netcdf\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1422\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1423\u001b[0m             \u001b[0;31m# No problems with the name - so we're fine!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1424\u001b[0;31m             \u001b[0mdataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1425\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1426\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_netcdf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Computing/miniconda3/envs/geo/lib/python3.6/site-packages/xarray/core/dataarray.py\u001b[0m in \u001b[0;36mto_dataset\u001b[0;34m(self, dim, name)\u001b[0m\n\u001b[1;32m    359\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_to_dataset_split\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    360\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 361\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_to_dataset_whole\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    362\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    363\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Computing/miniconda3/envs/geo/lib/python3.6/site-packages/xarray/core/dataarray.py\u001b[0m in \u001b[0;36m_to_dataset_whole\u001b[0;34m(self, name, shallow_copy)\u001b[0m\n\u001b[1;32m    326\u001b[0m                 \u001b[0mvariables\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvariables\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdeep\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    327\u001b[0m         \u001b[0mcoord_names\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_coords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 328\u001b[0;31m         \u001b[0mdataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_from_vars_and_coord_names\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvariables\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcoord_names\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    329\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mdataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    330\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Computing/miniconda3/envs/geo/lib/python3.6/site-packages/xarray/core/dataset.py\u001b[0m in \u001b[0;36m_from_vars_and_coord_names\u001b[0;34m(cls, variables, coord_names, attrs)\u001b[0m\n\u001b[1;32m    651\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mclassmethod\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    652\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_from_vars_and_coord_names\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcls\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvariables\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcoord_names\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattrs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 653\u001b[0;31m         \u001b[0mdims\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcalculate_dimensions\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvariables\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    654\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcls\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_construct_direct\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvariables\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcoord_names\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdims\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattrs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    655\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Research/Computing/miniconda3/envs/geo/lib/python3.6/site-packages/xarray/core/dataset.py\u001b[0m in \u001b[0;36mcalculate_dimensions\u001b[0;34m(variables)\u001b[0m\n\u001b[1;32m    111\u001b[0m                 raise ValueError('conflicting sizes for dimension %r: '\n\u001b[1;32m    112\u001b[0m                                  \u001b[0;34m'length %s on %r and length %s on %r'\u001b[0m \u001b[0;34m%\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m                                  (dim, size, k, dims[dim], last_used[dim]))\n\u001b[0m\u001b[1;32m    114\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mdims\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: conflicting sizes for dimension 'x': length 6 on 'data' and length 4 on 'x'"
     ]
    }
   ],
   "source": [
    "# inconsistent coordinate size is not allowed in a NetCDF file\n",
    "dr_out.to_netcdf('apply_ufunc_result.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This I/O issue can be fixed by correcting the coordinate variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dr_out['x'] = np.arange(6)\n",
    "dr_out.to_netcdf('apply_ufunc_result.nc') # no error now"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 2: Specifying a new output dimension name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we can give `output_core_dims` a different name. The new dimension `x_new` will have no coordinate value by default. This is still a legal NetCDF file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dr_out2 = xr.apply_ufunc(apply_A, dr, \n",
    "                         input_core_dims=[['x']],\n",
    "                         output_core_dims=[['x_new']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "dr_out2.to_netcdf('apply_ufunc_result.nc') # missing coordinate value is fine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 3: Specifying the same output dimension name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `output_core_dims` is the same as `input_core_dims`, the old, incorrect coordinate (`x` here) will be kept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (time: 2, x: 6)>\n",
       "array([[ 1.24765 ,  0.24795 ,  1.053913,  1.338334,  1.9124  ,  0.860453],\n",
       "       [ 0.923691,  0.432066,  0.751094,  1.274393,  1.106883,  0.899529]])\n",
       "Coordinates:\n",
       "  * time     (time) int64 0 1\n",
       "  * x        (x) int64 0 1 2 3"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.apply_ufunc(apply_A, dr, \n",
    "               input_core_dims=[['x']],\n",
    "               output_core_dims=[['x']],)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`xr.apply_ufunc` should probably drop the core dimension coordinate (`x` here) by default. This is safer (will not produce illegal NetCDF file) and also makes more physical sense (The core coordinate is likely to change)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 4: Specifying output coordinate size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`output_sizes` has no effect on either variable size or coordinate size, as long as the input is not a dask array..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.DataArray 'data' (time: 2, x: 6)>\n",
       "array([[ 1.24765 ,  0.24795 ,  1.053913,  1.338334,  1.9124  ,  0.860453],\n",
       "       [ 0.923691,  0.432066,  0.751094,  1.274393,  1.106883,  0.899529]])\n",
       "Coordinates:\n",
       "  * time     (time) int64 0 1\n",
       "  * x        (x) int64 0 1 2 3"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.apply_ufunc(apply_A, dr, \n",
    "               input_core_dims=[['x']],\n",
    "               output_core_dims=[['x']],\n",
    "               output_sizes = {'x': 999})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}


# Fixed code
def apply_variable_ufunc(func, *args, **kwargs):
    """apply_variable_ufunc(func, *args, signature, exclude_dims=frozenset())"""
    from .variable import Variable, as_compatible_data

    signature = kwargs.pop("signature")
    exclude_dims = kwargs.pop("exclude_dims", _DEFAULT_FROZEN_SET)
    dask = kwargs.pop("dask", "forbidden")
    output_dtypes = kwargs.pop("output_dtypes", None)
    output_sizes = kwargs.pop("output_sizes", None)
    keep_attrs = kwargs.pop("keep_attrs", False)
    if kwargs:
        raise TypeError(
            "apply_variable_ufunc() got unexpected keyword arguments: %s" % list(kwargs)
        )

    dim_sizes = unified_dim_sizes(
        (a for a in args if hasattr(a, "dims")), exclude_dims=exclude_dims
    )
    broadcast_dims = tuple(
        dim for dim in dim_sizes if dim not in signature.all_core_dims
    )
    output_dims = [broadcast_dims + out for out in signature.output_core_dims]

    input_data = [
        broadcast_compat_data(arg, broadcast_dims, core_dims)
        if isinstance(arg, Variable)
        else arg
        for arg, core_dims in zip(args, signature.input_core_dims)
    ]

    if any(isinstance(array, dask_array_type) for array in input_data):
        if dask == "forbidden":
            raise ValueError(
                "apply_ufunc encountered a dask array on an "
                "argument, but handling for dask arrays has not "
                "been enabled. Either set the ``dask`` argument "
                "or load your data into memory first with "
                "``.load()`` or ``.compute()``"
            )
        elif dask == "parallelized":
            input_dims = [broadcast_dims + dims for dims in signature.input_core_dims]
            numpy_func = func

            def func(*arrays):
                return _apply_with_dask_atop(
                    numpy_func,
                    arrays,
                    input_dims,
                    output_dims,
                    signature,
                    output_dtypes,
                    output_sizes,
                )
        elif dask == "allowed":
            pass
        else:
            raise ValueError(
                "unknown setting for dask array handling in apply_ufunc: {}".format(
                    dask
                )
            )
    result_data = func(*input_data)

    if signature.num_outputs == 1:
        result_data = (result_data,)
    elif (
        not isinstance(result_data, tuple) or len(result_data) != signature.num_outputs
    ):
        raise ValueError(
            "applied function does not have the number of "
            "outputs specified in the ufunc signature. "
            "Result is not a tuple of {} elements: {!r}".format(
                signature.num_outputs, result_data
            )
        )

    output = []
    for dims, data in zip(output_dims, result_data):
        data = as_compatible_data(data)
        if data.ndim != len(dims):
            raise ValueError(
                "applied function returned data with unexpected "
                "number of dimensions: {} vs {}, for dimensions {}".format(
                    data.ndim, len(dims), dims
                )
            )

        var = Variable(dims, data, fastpath=True)
        for dim, new_size in var.sizes.items():
            if dim in dim_sizes and new_size != dim_sizes[dim]:
                raise ValueError(
                    "size of dimension {!r} on inputs was unexpectedly "
                    "changed by applied function from {} to {}. Only "
                    "dimensions specified in ``exclude_dims`` with "
                    "xarray.apply_ufunc are allowed to change size.".format(
                        dim, dim_sizes[dim], new_size
                    )
                )

        if keep_attrs and isinstance(args[0], Variable):
            var.attrs.update(args[0].attrs)
        output.append(var)

    if signature.num_outputs == 1:
        return output[0]
    else:
        return tuple(output)