Merge pull request #29 from nschloe/outer

add npx.outer

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = npx
-version = 0.0.24
+version = 0.0.25
 author = Nico Schlömer
 author_email = nico.schloemer@gmail.com
 description = Some useful extensions for NumPy

src/npx/__init__.py

@@ -1,10 +1,11 @@
 from ._isin import isin_rows
-from ._main import add_at, dot, solve, subtract_at, sum_at
+from ._main import add_at, dot, outer, solve, subtract_at, sum_at
 from ._mean import mean
 from ._unique import unique, unique_rows
 
 __all__ = [
     "dot",
+    "outer",
     "solve",
     "sum_at",
     "add_at",

src/npx/_main.py

@@ -20,6 +20,15 @@ def dot(a: ArrayLike, b: ArrayLike) -> np.ndarray:
     return np.dot(a, b.reshape(b.shape[0], -1)).reshape(a.shape[:-1] + b.shape[1:])
 
 
+def outer(a: ArrayLike, b: ArrayLike) -> np.ndarray:
+    """Compute the outer product of two arrays `a` and `b` such that the shape
+    of the resulting array is `(*a.shape, *b.shape)`.
+    """
+    a = np.asarray(a)
+    b = np.asarray(b)
+    return np.outer(a, b).reshape(*a.shape, *b.shape)
+
+
 def solve(A: ArrayLike, x: ArrayLike) -> np.ndarray:
     """Solves a linear equation system with a matrix of shape (n, n) and an array of
     shape (n, ...). The output has the same shape as the second argument.

tests/test_dot_solve.py

@@ -17,3 +17,10 @@ def test_solve():
     b = np.random.rand(3, 4, 5)
     c = npx.solve(a, b)
     assert c.shape == b.shape
+
+
+def test_outer():
+    a = np.random.rand(1, 2)
+    b = np.random.rand(3, 4)
+    c = npx.outer(a, b)
+    assert c.shape == (1, 2, 3, 4)