implements expand_dims

expand_dims has the potential to hold more metadata than a plain reshape
and has cleaner axis-type propagation semantics

Author: kaushikcfd

pytato/__init__.py

@@ -39,6 +39,7 @@
         einsum,
 
         matmul, roll, transpose, stack, reshape, concatenate,
+        expand_dims,
 
         maximum, minimum, where,
 
@@ -95,7 +96,8 @@
         "make_dict_of_named_arrays", "make_placeholder", "make_size_param",
         "make_data_wrapper", "einsum",
 
-        "matmul", "roll", "transpose", "stack", "reshape", "concatenate",
+        "matmul", "roll", "transpose", "stack", "reshape", "expand_dims",
+        "concatenate",
 
         "generate_loopy", "generate_jax",
 

pytato/array.py

@@ -2464,4 +2464,44 @@ def squeeze(array: Array) -> Array:
             for i, s_i in enumerate(array.shape))]
 
 
+def expand_dims(array: Array, axis: Union[Tuple[int, ...], int]) -> Array:
+    """
+    Reshapes *array* by adding 1-long axes at *axis* dimensions of the returned
+    array.
+    """
+    from pytato.tags import ExpandedDimsReshape
+
+    if isinstance(axis, int):
+        axis = axis,
+
+    output_ndim = array.ndim + len(axis)
+
+    normalized_axis: List[int] = []
+
+    # {{{ sanity checks
+
+    for ax in axis:
+        if not (-output_ndim <= ax < output_ndim):
+            raise ValueError(f"Dimension {ax} not present in {output_ndim}-D array.")
+
+        normalized_axis.append(ax if ax >= 0 else (ax+output_ndim))
+
+    if len(set(normalized_axis)) != len(normalized_axis):
+        raise ValueError(f"repeated axis in '{axis}'.")
+
+    # }}}
+
+    new_shape = list(array.shape)
+
+    for ax in sorted(normalized_axis):
+        assert (0 <= ax < output_ndim)
+        new_shape.insert(ax, 1)
+
+    assert len(new_shape) == output_ndim
+
+    return Reshape(array, tuple(new_shape), "C",
+                   tags=(_get_default_tags()
+                         | {ExpandedDimsReshape(tuple(normalized_axis))}),
+                   axes=_get_default_axes(len(new_shape)))
+
 # vim: foldmethod=marker

pytato/tags.py

@@ -11,8 +11,9 @@
 .. autoclass:: AssumeNonNegative
 """
 
-
+from typing import Tuple
 from pytools.tag import Tag, UniqueTag, tag_dataclass
+from dataclasses import dataclass
 
 
 # {{{ pre-defined tag: ImplementationStrategy
@@ -101,3 +102,25 @@ class AssumeNonNegative(Tag):
     :class:`~pytato.target.Target` that all entries of the tagged array are
     non-negative.
     """
+
+
+@dataclass(eq=True, frozen=True, repr=True)
+class ExpandedDimsReshape(UniqueTag):
+    """
+    A tag that can be attached to a :class:`~pytato.array.Reshape` to indicate
+    that the new dimensions created by :func:`pytato.expand_dims`.
+
+    :attr new_dims: A :class:`tuple` of the dimensions of the reshaped array
+        that were added.
+
+    .. testsetup::
+
+        >>> import pytato as pt
+
+    .. doctest::
+
+        >>> x = pt.make_placeholder("x", (10, 4), "float64")
+        >>> pt.expand_dims(x, (0, 2, 4)).tags_of_type(pt.tags.ExpandedDimsReshape)
+        frozenset({ExpandedDimsReshape(new_dims=(0, 2, 4))})
+    """
+    new_dims: Tuple[int, ...]

test/test_codegen.py

@@ -1692,6 +1692,35 @@ def test_no_computation_for_empty_arrays(ctx_factory):
     assert not bprg.program.default_entrypoint.instructions
 
 
+def test_expand_dims(ctx_factory):
+    from numpy.random import default_rng
+
+    ntests = 50
+    rng = default_rng(seed=0)
+    ctx = ctx_factory()
+    cq = cl.CommandQueue(ctx)
+
+    for _ in range(ntests):
+        in_dim = rng.integers(2, 7)
+        n_new_axes = rng.integers(0, 7)
+        in_shape = rng.integers(2, 7, in_dim)
+        np_input = rng.random(in_shape, "float32")
+        pt_input = pt.make_data_wrapper(np_input)
+        axis = tuple(int(rng.integers(-(n_new_axes+in_dim), n_new_axes+in_dim))
+                     for _ in range(n_new_axes))
+
+        try:
+            np_output = np.expand_dims(np_input, axis=axis)
+        except ValueError:
+            with pytest.raises(ValueError):
+                pt.expand_dims(pt_input, axis=axis)
+        else:
+            _, (pt_output, ) = pt.generate_loopy(pt.expand_dims(pt_input,
+                                                                axis=axis))(cq)
+
+            np.testing.assert_allclose(np_output, pt_output)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])

test/test_pytato.py

@@ -884,6 +884,23 @@ def test_adv_indexing_into_zero_long_axes():
     # }}}
 
 
+def test_expand_dims_input_validate():
+    a = pt.make_placeholder("x", (10, 4), dtype="float64")
+
+    assert pt.expand_dims(a, (0, 2, 4)).shape == (1, 10, 1, 4, 1)
+    assert pt.expand_dims(a, (-5, -3, -1)).shape == (1, 10, 1, 4, 1)
+    assert pt.expand_dims(a, (-3)).shape == (1, 10, 4)
+
+    with pytest.raises(ValueError):
+        pt.expand_dims(a, (3, 3))
+
+    with pytest.raises(ValueError):
+        pt.expand_dims(a, (0, 2, 5))
+
+    with pytest.raises(ValueError):
+        pt.expand_dims(a, -4)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])