Support output_sizes in apply_gufunc

cubed/core/gufunc.py

@@ -10,6 +10,7 @@ def apply_gufunc(
     axes=None,
     axis=None,
     output_dtypes=None,
+    output_sizes=None,
     vectorize=None,
     **kwargs,
 ):
@@ -20,7 +21,7 @@ def apply_gufunc(
     `equivalent function <https://docs.dask.org/en/stable/generated/dask.array.gufunc.apply_gufunc.html>`_
     in Dask. Refer there for usage information.
 
-    Current limitations: ``keepdims``, ``output_sizes``, and ``allow_rechunk`` are not supported;
+    Current limitations: ``keepdims``, and ``allow_rechunk`` are not supported;
     and multiple outputs are not supported.
 
     Cubed assumes that ``func`` will allocate a new output array. However, if it allocates more memory
@@ -30,7 +31,6 @@ def apply_gufunc(
 
     # Currently the following parameters cannot be changed
     keepdims = False
-    output_sizes = None
     allow_rechunk = False
 
     # based on dask's apply_gufunc
@@ -142,4 +142,6 @@ def apply_gufunc(
     # Note (cubed): use blockwise on all output dimensions, not just loop_output_dims like in original
     out_ind = loop_output_dims + output_coredimss
 
-    return blockwise(func, out_ind, *arginds, dtype=output_dtypes, **kwargs)
+    return blockwise(
+        func, out_ind, *arginds, dtype=output_dtypes, new_axes=output_sizes, **kwargs
+    )

cubed/tests/test_gufunc.py

@@ -80,3 +80,24 @@ def outer_product(x, y):
 
     c = apply_gufunc(outer_product, "(i),(j)->(i,j)", a, b, vectorize=True)
     assert c.compute().shape == (10, 20, 30, 40)
+
+
+def test_gufunc_output_sizes(spec):
+    def foo(x):
+        return np.broadcast_to(x[:, np.newaxis], (x.shape[0], 3))
+
+    a = cubed.from_array(np.array([1, 2, 3, 4, 5], dtype=int), spec=spec)
+    x = apply_gufunc(foo, "()->(i_0)", a, output_dtypes=int, output_sizes={"i_0": 3})
+    assert x.chunks == ((5,), (3,))
+    assert_equal(
+        x,
+        np.array(
+            [
+                [1, 1, 1],
+                [2, 2, 2],
+                [3, 3, 3],
+                [4, 4, 4],
+                [5, 5, 5],
+            ]
+        ),
+    )