Generated docstring format + install from GitHub URL

README.md

@@ -26,11 +26,13 @@ print(psi.source)
 # from array_api_compat import array_namespace
 #
 # def psi_vec(x):
-#     """@vectorize
+#     """
+#     @vectorize
 #     def psi(x):
 #         if x < 0:
 #             return 0.0
-#         return x * math.exp(-x)"""
+#         return x * math.exp(-x)
+#     """
 #     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
 #     return xp.where(x < 0, 0.0, x * xp.exp(-x))
 ```
@@ -41,11 +43,13 @@ The same `psi` runs unchanged on torch tensors, JAX arrays, CuPy arrays, or
 ## Installation
 
 ```bash
-pip install -e ".[dev]"   # from a checkout
-make check                # lint + type + test + coverage
+pip install "array-vectorize @ git+https://github.com/grAItools/array-vectorize.git"
 ```
 
 Requires Python >= 3.12. The only runtime dependency is `array-api-compat`.
+The package is not on PyPI yet; once published, plain
+`pip install array-vectorize` will work. From a checkout (development):
+`pip install -e ".[dev]"` then `make check` (lint + type + test + coverage).
 
 ## Public API
 

docs/index.md

@@ -30,11 +30,13 @@ inspectable.
     from array_api_compat import array_namespace
 
     def psi_vec(x):
-        """@vectorize
+        """
+        @vectorize
         def psi(x):
             if x < 0:
                 return 0.0
-            return x * math.exp(-x)"""
+            return x * math.exp(-x)
+        """
         xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
         return xp.where(x < 0, 0.0, x * xp.exp(-x))
     ```
@@ -50,11 +52,13 @@ or `array_api_strict` arrays.
 ## Installation
 
 ```bash
-pip install -e ".[dev]"   # from a checkout
-make check                # lint + type + test + coverage
+pip install "array-vectorize @ git+https://github.com/grAItools/array-vectorize.git"
 ```
 
 Requires Python >= 3.12. The only runtime dependency is `array-api-compat`.
+The package is not on PyPI yet; once published, plain
+`pip install array-vectorize` will work. From a checkout (development):
+`pip install -e ".[dev]"` then `make check` (lint + type + test + coverage).
 
 ## Highlights
 

src/array_vectorize/codegen/emitter.py

@@ -258,8 +258,18 @@ def generate_source(lowered: LoweredFunction, program: Program, *, pinned: bool
         )
     else:
         ns_line = _namespace_line(ns, all_params)
+    # original source, verbatim — the docstring is cleandoc-exact:
+    # ast.get_docstring / inspect.getdoc on the generated function recover
+    # the original scalar source verbatim
+    indent = "    "  # docstring is always the first statement, body level
+    doc_value = (
+        "\n"
+        + "\n".join(indent + line if line.strip() else "" for line in lowered.source.splitlines())
+        + "\n"
+        + indent
+    )
     body: list[ast.stmt] = [
-        ast.Expr(value=ast.Constant(value=lowered.source)),  # original source, verbatim
+        ast.Expr(value=ast.Constant(value=doc_value)),
         ns_line,
     ]
     for stmt in program.bindings:

tests/golden/cases/abs_round.py

@@ -1,8 +1,10 @@
 from array_api_compat import array_namespace
 
 def abs_round_vec(x):
-    """def abs_round(x):
-    return abs(x) + round(x)"""
+    """
+    def abs_round(x):
+        return abs(x) + round(x)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     t_1 = xp.asarray(x)
     return xp.abs(t_1) + xp.round(t_1)

tests/golden/cases/add.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def add_vec(x, y):
-    """def add(x, y):
-    return x + y"""
+    """
+    def add(x, y):
+        return x + y
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return x + y

tests/golden/cases/add_pinned.py

@@ -1,5 +1,7 @@
 def add_vec(x, y, *, _namespace=None):
-    """def add(x, y):
-    return x + y"""
+    """
+    def add(x, y):
+        return x + y
+    """
     xp = _namespace
     return x + y

tests/golden/cases/arith_ops.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def arith_ops_vec(x, y):
-    """def arith_ops(x, y):
-    return (x + y - x * y) / (x * x + y * y)"""
+    """
+    def arith_ops(x, y):
+        return (x + y - x * y) / (x * x + y * y)
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return (x + y - x * y) / (x * x + y * y)

tests/golden/cases/bitwise_ops.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def bitwise_ops_vec(x, y):
-    """def bitwise_ops(x, y):
-    return (x & y) | (x ^ 3) << 1 >> 2"""
+    """
+    def bitwise_ops(x, y):
+        return (x & y) | (x ^ 3) << 1 >> 2
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return x & y | (x ^ 3) << 1 >> 2

tests/golden/cases/boolop_logical.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def boolop_logical_vec(x, y):
-    """def boolop_logical(x, y):
-    return ((x > 0) and (y > 0)) or not (x < 0)"""
+    """
+    def boolop_logical(x, y):
+        return ((x > 0) and (y > 0)) or not (x < 0)
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return xp.logical_or(xp.logical_and(x > 0, y > 0), xp.logical_not(x < 0))

tests/golden/cases/boolop_numeric.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def boolop_numeric_vec(x, y):
-    """def boolop_numeric(x, y):
-    return x and y"""
+    """
+    def boolop_numeric(x, y):
+        return x and y
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return xp.where(x != 0, y, x)

tests/golden/cases/both_branches_assign.py

@@ -1,14 +1,16 @@
 from array_api_compat import array_namespace
 
 def both_branches_assign_vec(x):
-    """def both_branches_assign(x):
-    if x > 0:
-        y = x
-        z = 1.0
-    else:
-        y = -x
-        z = 2.0
-    return y + z"""
+    """
+    def both_branches_assign(x):
+        if x > 0:
+            y = x
+            z = 1.0
+        else:
+            y = -x
+            z = 2.0
+        return y + z
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     y_1 = x
     z_1 = 1.0

tests/golden/cases/casts.py

@@ -1,8 +1,10 @@
 from array_api_compat import array_namespace
 
 def casts_vec(x):
-    """def casts(x):
-    return int(x) + float(x) + bool(x) + math.trunc(x)"""
+    """
+    def casts(x):
+        return int(x) + float(x) + bool(x) + math.trunc(x)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.int64)

tests/golden/cases/clamp.py

@@ -1,11 +1,13 @@
 from array_api_compat import array_namespace
 
 def clamp_vec(x, lo=0.0, hi=1.0):
-    """def clamp(x, lo=0.0, hi=1.0):
-    if x < lo:
-        return lo
-    if x > hi:
-        return hi
-    return x"""
+    """
+    def clamp(x, lo=0.0, hi=1.0):
+        if x < lo:
+            return lo
+        if x > hi:
+            return hi
+        return x
+    """
     xp = array_namespace(*[a for a in (x, lo, hi) if hasattr(a, '__array_namespace__')])
     return xp.where(x < lo, lo, xp.where(x > hi, hi, x))

tests/golden/cases/closure_array.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def closure_array_vec(x, *, ARR=None):
-    """def closure_array(x):
-    return x + ARR"""
+    """
+    def closure_array(x):
+        return x + ARR
+    """
     xp = array_namespace(*[a for a in (x, ARR) if hasattr(a, '__array_namespace__')])
     return x + ARR

tests/golden/cases/closure_scalar.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def closure_scalar_vec(x):
-    """def closure_scalar(x):
-    return x * SCALE"""
+    """
+    def closure_scalar(x):
+        return x * SCALE
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return x * 3.0

tests/golden/cases/compare_chain.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def compare_chain_vec(x):
-    """def compare_chain(x):
-    return 0 < x < 1"""
+    """
+    def compare_chain(x):
+        return 0 < x < 1
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return xp.logical_and(0 < x, x < 1)

tests/golden/cases/compare_ops.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def compare_ops_vec(x, y):
-    """def compare_ops(x, y):
-    return (x == y) + (x != y) + (x < y) + (x <= y) + (x > y) + (x >= y)"""
+    """
+    def compare_ops(x, y):
+        return (x == y) + (x != y) + (x < y) + (x <= y) + (x > y) + (x >= y)
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     return _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, _vec_arith(xp, 1, x == y, x != y), x < y), x <= y), x > y), x >= y)

tests/golden/cases/complex_expr.py

@@ -1,9 +1,11 @@
 from array_api_compat import array_namespace
 
 def complex_expr_vec(x, y):
-    """def complex_expr(x, y):
-    t = math.exp(-(x * x + y * y))
-    return t / (1.0 + t)"""
+    """
+    def complex_expr(x, y):
+        t = math.exp(-(x * x + y * y))
+        return t / (1.0 + t)
+    """
     xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')])
     t = xp.exp(xp.astype(xp.asarray(-(x * x + y * y)), xp.float64))
     return t / (1.0 + t)

tests/golden/cases/constants.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def constants_vec(x):
-    """def constants(x):
-    return x * math.pi + math.e + math.inf - math.nan"""
+    """
+    def constants(x):
+        return x * math.pi + math.e + math.inf - math.nan
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return x * 3.141592653589793 + 2.718281828459045 + xp.inf - xp.nan

tests/golden/cases/cse_opportunity.py

@@ -1,8 +1,10 @@
 from array_api_compat import array_namespace
 
 def cse_opportunity_vec(x):
-    """def cse_opportunity(x):
-    return math.sqrt(x) * math.sqrt(x) + math.sqrt(x)"""
+    """
+    def cse_opportunity(x):
+        return math.sqrt(x) * math.sqrt(x) + math.sqrt(x)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     t_1 = xp.asarray(x)
     t_2 = xp.astype(t_1, xp.float64)

tests/golden/cases/floor_mod_ops.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def floor_mod_ops_vec(x):
-    """def floor_mod_ops(x):
-    return (x // 2) + (x % 3)"""
+    """
+    def floor_mod_ops(x):
+        return (x // 2) + (x % 3)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return xp.floor_divide(x, 2) + xp.remainder(x, 3)

tests/golden/cases/from_import.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def from_import_vec(x):
-    """def from_import(x):
-    return fexp(x)"""
+    """
+    def from_import(x):
+        return fexp(x)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return xp.exp(xp.astype(xp.asarray(x), xp.float64))

tests/golden/cases/helper_outer.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def helper_outer_vec(x):
-    """def helper_outer(x):
-    return helper_inner(x) + helper_inner(x * 2.0)"""
+    """
+    def helper_outer(x):
+        return helper_inner(x) + helper_inner(x * 2.0)
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     return helper_inner_vec(x) + helper_inner_vec(x * 2.0)

tests/golden/cases/kwonly_defaults.py

@@ -1,7 +1,9 @@
 from array_api_compat import array_namespace
 
 def kwonly_defaults_vec(x, *, scale=2.0, bias=0.5):
-    """def kwonly_defaults(x, *, scale=2.0, bias=0.5):
-    return x * scale + bias"""
+    """
+    def kwonly_defaults(x, *, scale=2.0, bias=0.5):
+        return x * scale + bias
+    """
     xp = array_namespace(*[a for a in (x, scale, bias) if hasattr(a, '__array_namespace__')])
     return x * scale + bias

tests/golden/cases/loop_accumulate.py

@@ -1,11 +1,13 @@
 from array_api_compat import array_namespace
 
 def loop_accumulate_vec(x):
-    """def loop_accumulate(x):
-    s = 0.0
-    for i in range(4):  # noqa: B007
-        s = s + x
-    return s"""
+    """
+    def loop_accumulate(x):
+        s = 0.0
+        for i in range(4):  # noqa: B007
+            s = s + x
+        return s
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     s = 0.0
     s_1 = s

tests/golden/cases/loop_helper_caller.py

@@ -1,11 +1,13 @@
 from array_api_compat import array_namespace
 
 def loop_helper_caller_vec(x):
-    """def loop_helper_caller(x):
-    s = 0.0
-    for i in range(3):  # noqa: B007
-        s = s + helper_inner(x)
-    return s"""
+    """
+    def loop_helper_caller(x):
+        s = 0.0
+        for i in range(3):  # noqa: B007
+            s = s + helper_inner(x)
+        return s
+    """
     xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')])
     s = 0.0
     s_1 = s