From 545a1740c38045a86d10939229cd6ccf78eb2910 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Mon, 5 Oct 2026 19:46:03 +0200 Subject: [PATCH 1/3] codegen: generated docstring quotes on their own lines --- README.md | 6 ++++-- docs/index.md | 6 ++++-- src/array_vectorize/codegen/emitter.py | 4 +++- tests/golden/cases/abs_round.py | 6 ++++-- tests/golden/cases/add.py | 6 ++++-- tests/golden/cases/add_pinned.py | 6 ++++-- tests/golden/cases/arith_ops.py | 6 ++++-- tests/golden/cases/bitwise_ops.py | 6 ++++-- tests/golden/cases/boolop_logical.py | 6 ++++-- tests/golden/cases/boolop_numeric.py | 6 ++++-- tests/golden/cases/both_branches_assign.py | 6 ++++-- tests/golden/cases/casts.py | 6 ++++-- tests/golden/cases/clamp.py | 6 ++++-- tests/golden/cases/closure_array.py | 6 ++++-- tests/golden/cases/closure_scalar.py | 6 ++++-- tests/golden/cases/compare_chain.py | 6 ++++-- tests/golden/cases/compare_ops.py | 6 ++++-- tests/golden/cases/complex_expr.py | 6 ++++-- tests/golden/cases/constants.py | 6 ++++-- tests/golden/cases/cse_opportunity.py | 6 ++++-- tests/golden/cases/floor_mod_ops.py | 6 ++++-- tests/golden/cases/from_import.py | 6 ++++-- tests/golden/cases/helper_outer.py | 6 ++++-- tests/golden/cases/kwonly_defaults.py | 6 ++++-- tests/golden/cases/loop_accumulate.py | 6 ++++-- tests/golden/cases/loop_helper_caller.py | 6 ++++-- tests/golden/cases/loop_start_stop_step.py | 6 ++++-- tests/golden/cases/loop_two_carried.py | 6 ++++-- tests/golden/cases/loop_with_branch.py | 6 ++++-- tests/golden/cases/math_calls.py | 6 ++++-- tests/golden/cases/math_floor_ceil.py | 6 ++++-- tests/golden/cases/math_two_arg.py | 6 ++++-- tests/golden/cases/merge_with_prior.py | 6 ++++-- tests/golden/cases/minmax.py | 6 ++++-- tests/golden/cases/nested_early_returns.py | 6 ++++-- tests/golden/cases/nested_loops.py | 6 ++++-- tests/golden/cases/piecewise.py | 6 ++++-- tests/golden/cases/pow_op.py | 6 ++++-- tests/golden/cases/psi.py | 6 ++++-- tests/golden/cases/relu.py | 6 ++++-- tests/golden/cases/relu_with_else.py | 6 ++++-- tests/golden/cases/reserved_param.py | 6 ++++-- tests/golden/cases/ssa_rebind.py | 6 ++++-- tests/golden/cases/ternary.py | 6 ++++-- tests/golden/cases/unary_ops.py | 6 ++++-- tests/unit/test_codegen.py | 2 +- 46 files changed, 180 insertions(+), 90 deletions(-) diff --git a/README.md b/README.md index 4fce034..7569dce 100644 --- a/README.md +++ b/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)) ``` diff --git a/docs/index.md b/docs/index.md index 44bbfcd..b6544c5 100644 --- a/docs/index.md +++ b/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)) ``` diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index 5152857..d0636f7 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/src/array_vectorize/codegen/emitter.py @@ -259,7 +259,9 @@ def generate_source(lowered: LoweredFunction, program: Program, *, pinned: bool else: ns_line = _namespace_line(ns, all_params) body: list[ast.stmt] = [ - ast.Expr(value=ast.Constant(value=lowered.source)), # original source, verbatim + # original source, verbatim; surrounding newlines put the triple + # quotes on their own lines + ast.Expr(value=ast.Constant(value=f"\n{lowered.source}\n")), ns_line, ] for stmt in program.bindings: diff --git a/tests/golden/cases/abs_round.py b/tests/golden/cases/abs_round.py index fb2d4d4..cfb7268 100644 --- a/tests/golden/cases/abs_round.py +++ b/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) diff --git a/tests/golden/cases/add.py b/tests/golden/cases/add.py index 55a788c..5beef0e 100644 --- a/tests/golden/cases/add.py +++ b/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 diff --git a/tests/golden/cases/add_pinned.py b/tests/golden/cases/add_pinned.py index 227dff5..e92f10e 100644 --- a/tests/golden/cases/add_pinned.py +++ b/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 diff --git a/tests/golden/cases/arith_ops.py b/tests/golden/cases/arith_ops.py index eba3713..5bce0a1 100644 --- a/tests/golden/cases/arith_ops.py +++ b/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) diff --git a/tests/golden/cases/bitwise_ops.py b/tests/golden/cases/bitwise_ops.py index 5cfcdab..6655189 100644 --- a/tests/golden/cases/bitwise_ops.py +++ b/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 diff --git a/tests/golden/cases/boolop_logical.py b/tests/golden/cases/boolop_logical.py index 7bf6446..6cf1927 100644 --- a/tests/golden/cases/boolop_logical.py +++ b/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)) diff --git a/tests/golden/cases/boolop_numeric.py b/tests/golden/cases/boolop_numeric.py index 7fbe529..06e8de0 100644 --- a/tests/golden/cases/boolop_numeric.py +++ b/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) diff --git a/tests/golden/cases/both_branches_assign.py b/tests/golden/cases/both_branches_assign.py index 63f06ee..43e77f6 100644 --- a/tests/golden/cases/both_branches_assign.py +++ b/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): + """ +def both_branches_assign(x): if x > 0: y = x z = 1.0 else: y = -x z = 2.0 - return y + z""" + return y + z +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) y_1 = x z_1 = 1.0 diff --git a/tests/golden/cases/casts.py b/tests/golden/cases/casts.py index aafce8c..e4bec16 100644 --- a/tests/golden/cases/casts.py +++ b/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) diff --git a/tests/golden/cases/clamp.py b/tests/golden/cases/clamp.py index f06cccc..25c526d 100644 --- a/tests/golden/cases/clamp.py +++ b/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): + """ +def clamp(x, lo=0.0, hi=1.0): if x < lo: return lo if x > hi: return hi - return x""" + 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)) diff --git a/tests/golden/cases/closure_array.py b/tests/golden/cases/closure_array.py index 758baa6..ea6c932 100644 --- a/tests/golden/cases/closure_array.py +++ b/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 diff --git a/tests/golden/cases/closure_scalar.py b/tests/golden/cases/closure_scalar.py index d335ce2..c86acf0 100644 --- a/tests/golden/cases/closure_scalar.py +++ b/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 diff --git a/tests/golden/cases/compare_chain.py b/tests/golden/cases/compare_chain.py index b817ed1..1fe8f4f 100644 --- a/tests/golden/cases/compare_chain.py +++ b/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) diff --git a/tests/golden/cases/compare_ops.py b/tests/golden/cases/compare_ops.py index d400f83..b845e0f 100644 --- a/tests/golden/cases/compare_ops.py +++ b/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) diff --git a/tests/golden/cases/complex_expr.py b/tests/golden/cases/complex_expr.py index 749a5df..82a48ad 100644 --- a/tests/golden/cases/complex_expr.py +++ b/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): + """ +def complex_expr(x, y): t = math.exp(-(x * x + y * y)) - return t / (1.0 + t)""" + 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) diff --git a/tests/golden/cases/constants.py b/tests/golden/cases/constants.py index cacdaa9..3084534 100644 --- a/tests/golden/cases/constants.py +++ b/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 diff --git a/tests/golden/cases/cse_opportunity.py b/tests/golden/cases/cse_opportunity.py index d2e6f73..99839e2 100644 --- a/tests/golden/cases/cse_opportunity.py +++ b/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) diff --git a/tests/golden/cases/floor_mod_ops.py b/tests/golden/cases/floor_mod_ops.py index 4c4c298..c2786b6 100644 --- a/tests/golden/cases/floor_mod_ops.py +++ b/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) diff --git a/tests/golden/cases/from_import.py b/tests/golden/cases/from_import.py index b51841a..da4086d 100644 --- a/tests/golden/cases/from_import.py +++ b/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)) diff --git a/tests/golden/cases/helper_outer.py b/tests/golden/cases/helper_outer.py index 8df8142..1e997df 100644 --- a/tests/golden/cases/helper_outer.py +++ b/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) diff --git a/tests/golden/cases/kwonly_defaults.py b/tests/golden/cases/kwonly_defaults.py index 067cd59..4523882 100644 --- a/tests/golden/cases/kwonly_defaults.py +++ b/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 diff --git a/tests/golden/cases/loop_accumulate.py b/tests/golden/cases/loop_accumulate.py index a946f96..d0380bf 100644 --- a/tests/golden/cases/loop_accumulate.py +++ b/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): + """ +def loop_accumulate(x): s = 0.0 for i in range(4): # noqa: B007 s = s + x - return s""" + return s +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 0.0 s_1 = s diff --git a/tests/golden/cases/loop_helper_caller.py b/tests/golden/cases/loop_helper_caller.py index 5ef8102..6bc9a4a 100644 --- a/tests/golden/cases/loop_helper_caller.py +++ b/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): + """ +def loop_helper_caller(x): s = 0.0 for i in range(3): # noqa: B007 s = s + helper_inner(x) - return s""" + return s +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 0.0 s_1 = s diff --git a/tests/golden/cases/loop_start_stop_step.py b/tests/golden/cases/loop_start_stop_step.py index 1b12638..485b343 100644 --- a/tests/golden/cases/loop_start_stop_step.py +++ b/tests/golden/cases/loop_start_stop_step.py @@ -1,11 +1,13 @@ from array_api_compat import array_namespace def loop_start_stop_step_vec(x): - """def loop_start_stop_step(x): + """ +def loop_start_stop_step(x): s = 1.0 for i in range(2, 8, 3): s = s * i - return s + x""" + return s + x +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 1.0 s_1 = s diff --git a/tests/golden/cases/loop_two_carried.py b/tests/golden/cases/loop_two_carried.py index db32fbe..1e4aebd 100644 --- a/tests/golden/cases/loop_two_carried.py +++ b/tests/golden/cases/loop_two_carried.py @@ -1,13 +1,15 @@ from array_api_compat import array_namespace def loop_two_carried_vec(x): - """def loop_two_carried(x): + """ +def loop_two_carried(x): a = 0.0 b = 1.0 for i in range(3): # noqa: B007 a = a + x b = b * 2.0 - return a + b""" + return a + b +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) a = 0.0 b = 1.0 diff --git a/tests/golden/cases/loop_with_branch.py b/tests/golden/cases/loop_with_branch.py index b186031..3e133b2 100644 --- a/tests/golden/cases/loop_with_branch.py +++ b/tests/golden/cases/loop_with_branch.py @@ -1,14 +1,16 @@ from array_api_compat import array_namespace def loop_with_branch_vec(x): - """def loop_with_branch(x): + """ +def loop_with_branch(x): s = 0.0 for i in range(4): # noqa: B007 if x > 0: # noqa: SIM108 s = s + x else: s = s - 1.0 - return s""" + return s +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 0.0 s_1 = s diff --git a/tests/golden/cases/math_calls.py b/tests/golden/cases/math_calls.py index 3216012..6baf7ee 100644 --- a/tests/golden/cases/math_calls.py +++ b/tests/golden/cases/math_calls.py @@ -1,8 +1,10 @@ from array_api_compat import array_namespace def math_calls_vec(x): - """def math_calls(x): - return math.sqrt(x) + math.exp(x) + math.log1p(x) + math.sin(x)""" + """ +def math_calls(x): + return math.sqrt(x) + math.exp(x) + math.log1p(x) + math.sin(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) diff --git a/tests/golden/cases/math_floor_ceil.py b/tests/golden/cases/math_floor_ceil.py index c262422..d6d6e72 100644 --- a/tests/golden/cases/math_floor_ceil.py +++ b/tests/golden/cases/math_floor_ceil.py @@ -1,8 +1,10 @@ from array_api_compat import array_namespace def math_floor_ceil_vec(x): - """def math_floor_ceil(x): - return math.floor(x) + math.ceil(x)""" + """ +def math_floor_ceil(x): + return math.floor(x) + math.ceil(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) diff --git a/tests/golden/cases/math_two_arg.py b/tests/golden/cases/math_two_arg.py index 081df37..8138e4a 100644 --- a/tests/golden/cases/math_two_arg.py +++ b/tests/golden/cases/math_two_arg.py @@ -1,8 +1,10 @@ from array_api_compat import array_namespace def math_two_arg_vec(x, y): - """def math_two_arg(x, y): - return math.atan2(x, y) + math.hypot(x, y) + math.copysign(x, y)""" + """ +def math_two_arg(x, y): + return math.atan2(x, y) + math.hypot(x, y) + math.copysign(x, y) +""" xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')]) t_1 = xp.asarray(x) t_2 = xp.asarray(y) diff --git a/tests/golden/cases/merge_with_prior.py b/tests/golden/cases/merge_with_prior.py index 1b4c310..dc6df4b 100644 --- a/tests/golden/cases/merge_with_prior.py +++ b/tests/golden/cases/merge_with_prior.py @@ -1,11 +1,13 @@ from array_api_compat import array_namespace def merge_with_prior_vec(x): - """def merge_with_prior(x): + """ +def merge_with_prior(x): r = 0.0 if x > 0: r = x - return r""" + return r +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) r = 0.0 r_1 = x diff --git a/tests/golden/cases/minmax.py b/tests/golden/cases/minmax.py index 4f28664..dd81663 100644 --- a/tests/golden/cases/minmax.py +++ b/tests/golden/cases/minmax.py @@ -1,8 +1,10 @@ from array_api_compat import array_namespace def minmax_vec(x, y): - """def minmax(x, y): - return min(x, y) + max(x, y, 3.0)""" + """ +def minmax(x, y): + return min(x, y) + max(x, y, 3.0) +""" xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')]) t_1 = xp.asarray(x) t_2 = xp.asarray(y) diff --git a/tests/golden/cases/nested_early_returns.py b/tests/golden/cases/nested_early_returns.py index 3ddc381..5e1681f 100644 --- a/tests/golden/cases/nested_early_returns.py +++ b/tests/golden/cases/nested_early_returns.py @@ -1,11 +1,13 @@ from array_api_compat import array_namespace def nested_early_returns_vec(x): - """def nested_early_returns(x): + """ +def nested_early_returns(x): if x > 0: if x > 10: return 1.0 return 2.0 - return 3.0""" + return 3.0 +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, xp.where(x > 10, 1.0, 2.0), 3.0) diff --git a/tests/golden/cases/nested_loops.py b/tests/golden/cases/nested_loops.py index c3d902a..6aa5fc7 100644 --- a/tests/golden/cases/nested_loops.py +++ b/tests/golden/cases/nested_loops.py @@ -1,12 +1,14 @@ from array_api_compat import array_namespace def nested_loops_vec(x): - """def nested_loops(x): + """ +def nested_loops(x): s = 0.0 for i in range(3): # noqa: B007 for j in range(2): # noqa: B007 # noqa: B007 s = s + x - return s""" + return s +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 0.0 s_1 = s diff --git a/tests/golden/cases/piecewise.py b/tests/golden/cases/piecewise.py index b415801..4fc0c48 100644 --- a/tests/golden/cases/piecewise.py +++ b/tests/golden/cases/piecewise.py @@ -1,11 +1,13 @@ from array_api_compat import array_namespace def piecewise_vec(x): - """def piecewise(x): + """ +def piecewise(x): if x < -1: return -1.0 elif x > 1: return 1.0 - return x""" + return x +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x < -1, -1.0, xp.where(x > 1, 1.0, x)) diff --git a/tests/golden/cases/pow_op.py b/tests/golden/cases/pow_op.py index 00db9a8..d7bde61 100644 --- a/tests/golden/cases/pow_op.py +++ b/tests/golden/cases/pow_op.py @@ -1,7 +1,9 @@ from array_api_compat import array_namespace def pow_op_vec(x): - """def pow_op(x): - return x**2 + 2.0**x""" + """ +def pow_op(x): + return x**2 + 2.0**x +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.pow(x, 2) + xp.pow(2.0, x) diff --git a/tests/golden/cases/psi.py b/tests/golden/cases/psi.py index c92c618..7269171 100644 --- a/tests/golden/cases/psi.py +++ b/tests/golden/cases/psi.py @@ -1,9 +1,11 @@ from array_api_compat import array_namespace def psi_vec(x): - """def psi(x): + """ +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(xp.astype(xp.asarray(-x), xp.float64))) diff --git a/tests/golden/cases/relu.py b/tests/golden/cases/relu.py index 1a139b2..7fb574a 100644 --- a/tests/golden/cases/relu.py +++ b/tests/golden/cases/relu.py @@ -1,9 +1,11 @@ from array_api_compat import array_namespace def relu_vec(x): - """def relu(x): + """ +def relu(x): if x > 0: return x - return 0.0""" + return 0.0 +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, x, 0.0) diff --git a/tests/golden/cases/relu_with_else.py b/tests/golden/cases/relu_with_else.py index fd29c2f..b3a8039 100644 --- a/tests/golden/cases/relu_with_else.py +++ b/tests/golden/cases/relu_with_else.py @@ -1,12 +1,14 @@ from array_api_compat import array_namespace def relu_with_else_vec(x): - """def relu_with_else(x): + """ +def relu_with_else(x): if x > 0: # noqa: SIM108 r = x else: r = 0.0 - return r""" + return r +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) r_1 = x r_2 = 0.0 diff --git a/tests/golden/cases/reserved_param.py b/tests/golden/cases/reserved_param.py index c790ca0..2f48fa3 100644 --- a/tests/golden/cases/reserved_param.py +++ b/tests/golden/cases/reserved_param.py @@ -1,7 +1,9 @@ from array_api_compat import array_namespace def reserved_param_vec(xp): - """def reserved_param(xp): - return xp + 1""" + """ +def reserved_param(xp): + return xp + 1 +""" xp_1 = array_namespace(*[a for a in (xp,) if hasattr(a, '__array_namespace__')]) return xp + 1 diff --git a/tests/golden/cases/ssa_rebind.py b/tests/golden/cases/ssa_rebind.py index 899d006..e49bc50 100644 --- a/tests/golden/cases/ssa_rebind.py +++ b/tests/golden/cases/ssa_rebind.py @@ -1,11 +1,13 @@ from array_api_compat import array_namespace def ssa_rebind_vec(x): - """def ssa_rebind(x): + """ +def ssa_rebind(x): y = x + 1 y = y * 2 y += 3 - return y""" + return y +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) y = x + 1 y_1 = y * 2 diff --git a/tests/golden/cases/ternary.py b/tests/golden/cases/ternary.py index a683791..4d1aa3f 100644 --- a/tests/golden/cases/ternary.py +++ b/tests/golden/cases/ternary.py @@ -1,7 +1,9 @@ from array_api_compat import array_namespace def ternary_vec(x): - """def ternary(x): - return x * 2 if x > 0 else x / 2""" + """ +def ternary(x): + return x * 2 if x > 0 else x / 2 +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, x * 2, x / 2) diff --git a/tests/golden/cases/unary_ops.py b/tests/golden/cases/unary_ops.py index 159a3a0..7dba06a 100644 --- a/tests/golden/cases/unary_ops.py +++ b/tests/golden/cases/unary_ops.py @@ -1,7 +1,9 @@ from array_api_compat import array_namespace def unary_ops_vec(x): - """def unary_ops(x): - return -x + +x - ~x""" + """ +def unary_ops(x): + return -x + +x - ~x +""" xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return -x + +x - ~x diff --git a/tests/unit/test_codegen.py b/tests/unit/test_codegen.py index 94a564f..8611754 100644 --- a/tests/unit/test_codegen.py +++ b/tests/unit/test_codegen.py @@ -51,7 +51,7 @@ def test_module_shape() -> None: src = gen(Program(("x",), (), Ref("x"))) assert src.startswith("from array_api_compat import array_namespace") assert "def f_vec(x):" in src - assert '"""def f(x):\n return x"""' in src + assert '"""\ndef f(x):\n return x\n"""' in src assert "xp = array_namespace(" in src assert "hasattr(a, '__array_namespace__')" in src From 65667c4cb954d07e9dfecb16370c2897592dee52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Mon, 5 Oct 2026 19:46:15 +0200 Subject: [PATCH 2/3] docs: install from the github repo url --- README.md | 6 ++++-- docs/index.md | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 7569dce..f2bc2e1 100644 --- a/README.md +++ b/README.md @@ -43,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 diff --git a/docs/index.md b/docs/index.md index b6544c5..b291c7a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -52,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 From a3c0292a6f5a417486b85ef7c560ce79d891555a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Mon, 5 Oct 2026 20:03:48 +0200 Subject: [PATCH 3/3] codegen: indent generated docstring content to body level --- src/array_vectorize/codegen/emitter.py | 14 +++++++++++--- tests/golden/cases/abs_round.py | 6 +++--- tests/golden/cases/add.py | 6 +++--- tests/golden/cases/add_pinned.py | 6 +++--- tests/golden/cases/arith_ops.py | 6 +++--- tests/golden/cases/bitwise_ops.py | 6 +++--- tests/golden/cases/boolop_logical.py | 6 +++--- tests/golden/cases/boolop_numeric.py | 6 +++--- tests/golden/cases/both_branches_assign.py | 18 +++++++++--------- tests/golden/cases/casts.py | 6 +++--- tests/golden/cases/clamp.py | 14 +++++++------- tests/golden/cases/closure_array.py | 6 +++--- tests/golden/cases/closure_scalar.py | 6 +++--- tests/golden/cases/compare_chain.py | 6 +++--- tests/golden/cases/compare_ops.py | 6 +++--- tests/golden/cases/complex_expr.py | 8 ++++---- tests/golden/cases/constants.py | 6 +++--- tests/golden/cases/cse_opportunity.py | 6 +++--- tests/golden/cases/floor_mod_ops.py | 6 +++--- tests/golden/cases/from_import.py | 6 +++--- tests/golden/cases/helper_outer.py | 6 +++--- tests/golden/cases/kwonly_defaults.py | 6 +++--- tests/golden/cases/loop_accumulate.py | 12 ++++++------ tests/golden/cases/loop_helper_caller.py | 12 ++++++------ tests/golden/cases/loop_start_stop_step.py | 12 ++++++------ tests/golden/cases/loop_two_carried.py | 16 ++++++++-------- tests/golden/cases/loop_with_branch.py | 18 +++++++++--------- tests/golden/cases/math_calls.py | 6 +++--- tests/golden/cases/math_floor_ceil.py | 6 +++--- tests/golden/cases/math_two_arg.py | 6 +++--- tests/golden/cases/merge_with_prior.py | 12 ++++++------ tests/golden/cases/minmax.py | 6 +++--- tests/golden/cases/nested_early_returns.py | 14 +++++++------- tests/golden/cases/nested_loops.py | 14 +++++++------- tests/golden/cases/piecewise.py | 14 +++++++------- tests/golden/cases/pow_op.py | 6 +++--- tests/golden/cases/psi.py | 10 +++++----- tests/golden/cases/relu.py | 10 +++++----- tests/golden/cases/relu_with_else.py | 14 +++++++------- tests/golden/cases/reserved_param.py | 6 +++--- tests/golden/cases/ssa_rebind.py | 12 ++++++------ tests/golden/cases/ternary.py | 6 +++--- tests/golden/cases/unary_ops.py | 6 +++--- tests/test_inspectability.py | 12 ++++++++++++ tests/unit/test_codegen.py | 2 +- 45 files changed, 207 insertions(+), 187 deletions(-) diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index d0636f7..3ef51ed 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/src/array_vectorize/codegen/emitter.py @@ -258,10 +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] = [ - # original source, verbatim; surrounding newlines put the triple - # quotes on their own lines - ast.Expr(value=ast.Constant(value=f"\n{lowered.source}\n")), + ast.Expr(value=ast.Constant(value=doc_value)), ns_line, ] for stmt in program.bindings: diff --git a/tests/golden/cases/abs_round.py b/tests/golden/cases/abs_round.py index cfb7268..d9a67d5 100644 --- a/tests/golden/cases/abs_round.py +++ b/tests/golden/cases/abs_round.py @@ -2,9 +2,9 @@ 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) diff --git a/tests/golden/cases/add.py b/tests/golden/cases/add.py index 5beef0e..80bffb9 100644 --- a/tests/golden/cases/add.py +++ b/tests/golden/cases/add.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/add_pinned.py b/tests/golden/cases/add_pinned.py index e92f10e..d697e8a 100644 --- a/tests/golden/cases/add_pinned.py +++ b/tests/golden/cases/add_pinned.py @@ -1,7 +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 diff --git a/tests/golden/cases/arith_ops.py b/tests/golden/cases/arith_ops.py index 5bce0a1..66e39a2 100644 --- a/tests/golden/cases/arith_ops.py +++ b/tests/golden/cases/arith_ops.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/bitwise_ops.py b/tests/golden/cases/bitwise_ops.py index 6655189..0406f21 100644 --- a/tests/golden/cases/bitwise_ops.py +++ b/tests/golden/cases/bitwise_ops.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/boolop_logical.py b/tests/golden/cases/boolop_logical.py index 6cf1927..b44f6f7 100644 --- a/tests/golden/cases/boolop_logical.py +++ b/tests/golden/cases/boolop_logical.py @@ -2,8 +2,8 @@ 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)) diff --git a/tests/golden/cases/boolop_numeric.py b/tests/golden/cases/boolop_numeric.py index 06e8de0..336077b 100644 --- a/tests/golden/cases/boolop_numeric.py +++ b/tests/golden/cases/boolop_numeric.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/both_branches_assign.py b/tests/golden/cases/both_branches_assign.py index 43e77f6..2a3ed20 100644 --- a/tests/golden/cases/both_branches_assign.py +++ b/tests/golden/cases/both_branches_assign.py @@ -2,15 +2,15 @@ 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 diff --git a/tests/golden/cases/casts.py b/tests/golden/cases/casts.py index e4bec16..1befbfe 100644 --- a/tests/golden/cases/casts.py +++ b/tests/golden/cases/casts.py @@ -2,9 +2,9 @@ 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) diff --git a/tests/golden/cases/clamp.py b/tests/golden/cases/clamp.py index 25c526d..904df34 100644 --- a/tests/golden/cases/clamp.py +++ b/tests/golden/cases/clamp.py @@ -2,12 +2,12 @@ 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)) diff --git a/tests/golden/cases/closure_array.py b/tests/golden/cases/closure_array.py index ea6c932..b657fd0 100644 --- a/tests/golden/cases/closure_array.py +++ b/tests/golden/cases/closure_array.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/closure_scalar.py b/tests/golden/cases/closure_scalar.py index c86acf0..321db23 100644 --- a/tests/golden/cases/closure_scalar.py +++ b/tests/golden/cases/closure_scalar.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/compare_chain.py b/tests/golden/cases/compare_chain.py index 1fe8f4f..702fdc4 100644 --- a/tests/golden/cases/compare_chain.py +++ b/tests/golden/cases/compare_chain.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/compare_ops.py b/tests/golden/cases/compare_ops.py index b845e0f..b5f643d 100644 --- a/tests/golden/cases/compare_ops.py +++ b/tests/golden/cases/compare_ops.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/complex_expr.py b/tests/golden/cases/complex_expr.py index 82a48ad..da9c8c3 100644 --- a/tests/golden/cases/complex_expr.py +++ b/tests/golden/cases/complex_expr.py @@ -2,10 +2,10 @@ 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) diff --git a/tests/golden/cases/constants.py b/tests/golden/cases/constants.py index 3084534..f6c4e49 100644 --- a/tests/golden/cases/constants.py +++ b/tests/golden/cases/constants.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/cse_opportunity.py b/tests/golden/cases/cse_opportunity.py index 99839e2..acffa8c 100644 --- a/tests/golden/cases/cse_opportunity.py +++ b/tests/golden/cases/cse_opportunity.py @@ -2,9 +2,9 @@ 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) diff --git a/tests/golden/cases/floor_mod_ops.py b/tests/golden/cases/floor_mod_ops.py index c2786b6..c1a0678 100644 --- a/tests/golden/cases/floor_mod_ops.py +++ b/tests/golden/cases/floor_mod_ops.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/from_import.py b/tests/golden/cases/from_import.py index da4086d..000b4d7 100644 --- a/tests/golden/cases/from_import.py +++ b/tests/golden/cases/from_import.py @@ -2,8 +2,8 @@ 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)) diff --git a/tests/golden/cases/helper_outer.py b/tests/golden/cases/helper_outer.py index 1e997df..078d866 100644 --- a/tests/golden/cases/helper_outer.py +++ b/tests/golden/cases/helper_outer.py @@ -2,8 +2,8 @@ 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) diff --git a/tests/golden/cases/kwonly_defaults.py b/tests/golden/cases/kwonly_defaults.py index 4523882..f7c76dd 100644 --- a/tests/golden/cases/kwonly_defaults.py +++ b/tests/golden/cases/kwonly_defaults.py @@ -2,8 +2,8 @@ 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 diff --git a/tests/golden/cases/loop_accumulate.py b/tests/golden/cases/loop_accumulate.py index d0380bf..13f886a 100644 --- a/tests/golden/cases/loop_accumulate.py +++ b/tests/golden/cases/loop_accumulate.py @@ -2,12 +2,12 @@ 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 diff --git a/tests/golden/cases/loop_helper_caller.py b/tests/golden/cases/loop_helper_caller.py index 6bc9a4a..666b40f 100644 --- a/tests/golden/cases/loop_helper_caller.py +++ b/tests/golden/cases/loop_helper_caller.py @@ -2,12 +2,12 @@ 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 diff --git a/tests/golden/cases/loop_start_stop_step.py b/tests/golden/cases/loop_start_stop_step.py index 485b343..f9a9461 100644 --- a/tests/golden/cases/loop_start_stop_step.py +++ b/tests/golden/cases/loop_start_stop_step.py @@ -2,12 +2,12 @@ def loop_start_stop_step_vec(x): """ -def loop_start_stop_step(x): - s = 1.0 - for i in range(2, 8, 3): - s = s * i - return s + x -""" + def loop_start_stop_step(x): + s = 1.0 + for i in range(2, 8, 3): + s = s * i + return s + x + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 1.0 s_1 = s diff --git a/tests/golden/cases/loop_two_carried.py b/tests/golden/cases/loop_two_carried.py index 1e4aebd..0bf172d 100644 --- a/tests/golden/cases/loop_two_carried.py +++ b/tests/golden/cases/loop_two_carried.py @@ -2,14 +2,14 @@ def loop_two_carried_vec(x): """ -def loop_two_carried(x): - a = 0.0 - b = 1.0 - for i in range(3): # noqa: B007 - a = a + x - b = b * 2.0 - return a + b -""" + def loop_two_carried(x): + a = 0.0 + b = 1.0 + for i in range(3): # noqa: B007 + a = a + x + b = b * 2.0 + return a + b + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) a = 0.0 b = 1.0 diff --git a/tests/golden/cases/loop_with_branch.py b/tests/golden/cases/loop_with_branch.py index 3e133b2..987b845 100644 --- a/tests/golden/cases/loop_with_branch.py +++ b/tests/golden/cases/loop_with_branch.py @@ -2,15 +2,15 @@ def loop_with_branch_vec(x): """ -def loop_with_branch(x): - s = 0.0 - for i in range(4): # noqa: B007 - if x > 0: # noqa: SIM108 - s = s + x - else: - s = s - 1.0 - return s -""" + def loop_with_branch(x): + s = 0.0 + for i in range(4): # noqa: B007 + if x > 0: # noqa: SIM108 + s = s + x + else: + s = s - 1.0 + return s + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) s = 0.0 s_1 = s diff --git a/tests/golden/cases/math_calls.py b/tests/golden/cases/math_calls.py index 6baf7ee..f0f04f2 100644 --- a/tests/golden/cases/math_calls.py +++ b/tests/golden/cases/math_calls.py @@ -2,9 +2,9 @@ def math_calls_vec(x): """ -def math_calls(x): - return math.sqrt(x) + math.exp(x) + math.log1p(x) + math.sin(x) -""" + def math_calls(x): + return math.sqrt(x) + math.exp(x) + math.log1p(x) + math.sin(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) diff --git a/tests/golden/cases/math_floor_ceil.py b/tests/golden/cases/math_floor_ceil.py index d6d6e72..495a20a 100644 --- a/tests/golden/cases/math_floor_ceil.py +++ b/tests/golden/cases/math_floor_ceil.py @@ -2,9 +2,9 @@ def math_floor_ceil_vec(x): """ -def math_floor_ceil(x): - return math.floor(x) + math.ceil(x) -""" + def math_floor_ceil(x): + return math.floor(x) + math.ceil(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) diff --git a/tests/golden/cases/math_two_arg.py b/tests/golden/cases/math_two_arg.py index 8138e4a..07c7397 100644 --- a/tests/golden/cases/math_two_arg.py +++ b/tests/golden/cases/math_two_arg.py @@ -2,9 +2,9 @@ def math_two_arg_vec(x, y): """ -def math_two_arg(x, y): - return math.atan2(x, y) + math.hypot(x, y) + math.copysign(x, y) -""" + def math_two_arg(x, y): + return math.atan2(x, y) + math.hypot(x, y) + math.copysign(x, y) + """ xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')]) t_1 = xp.asarray(x) t_2 = xp.asarray(y) diff --git a/tests/golden/cases/merge_with_prior.py b/tests/golden/cases/merge_with_prior.py index dc6df4b..90992f0 100644 --- a/tests/golden/cases/merge_with_prior.py +++ b/tests/golden/cases/merge_with_prior.py @@ -2,12 +2,12 @@ def merge_with_prior_vec(x): """ -def merge_with_prior(x): - r = 0.0 - if x > 0: - r = x - return r -""" + def merge_with_prior(x): + r = 0.0 + if x > 0: + r = x + return r + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) r = 0.0 r_1 = x diff --git a/tests/golden/cases/minmax.py b/tests/golden/cases/minmax.py index dd81663..27af9fb 100644 --- a/tests/golden/cases/minmax.py +++ b/tests/golden/cases/minmax.py @@ -2,9 +2,9 @@ def minmax_vec(x, y): """ -def minmax(x, y): - return min(x, y) + max(x, y, 3.0) -""" + def minmax(x, y): + return min(x, y) + max(x, y, 3.0) + """ xp = array_namespace(*[a for a in (x, y) if hasattr(a, '__array_namespace__')]) t_1 = xp.asarray(x) t_2 = xp.asarray(y) diff --git a/tests/golden/cases/nested_early_returns.py b/tests/golden/cases/nested_early_returns.py index 5e1681f..c8e7a61 100644 --- a/tests/golden/cases/nested_early_returns.py +++ b/tests/golden/cases/nested_early_returns.py @@ -2,12 +2,12 @@ def nested_early_returns_vec(x): """ -def nested_early_returns(x): - if x > 0: - if x > 10: - return 1.0 - return 2.0 - return 3.0 -""" + def nested_early_returns(x): + if x > 0: + if x > 10: + return 1.0 + return 2.0 + return 3.0 + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, xp.where(x > 10, 1.0, 2.0), 3.0) diff --git a/tests/golden/cases/nested_loops.py b/tests/golden/cases/nested_loops.py index 6aa5fc7..3092987 100644 --- a/tests/golden/cases/nested_loops.py +++ b/tests/golden/cases/nested_loops.py @@ -2,13 +2,13 @@ def nested_loops_vec(x): """ -def nested_loops(x): - s = 0.0 - for i in range(3): # noqa: B007 - for j in range(2): # noqa: B007 # noqa: B007 - s = s + x - return s -""" + def nested_loops(x): + s = 0.0 + for i in range(3): # noqa: B007 + for j in range(2): # noqa: B007 # 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 diff --git a/tests/golden/cases/piecewise.py b/tests/golden/cases/piecewise.py index 4fc0c48..153885f 100644 --- a/tests/golden/cases/piecewise.py +++ b/tests/golden/cases/piecewise.py @@ -2,12 +2,12 @@ def piecewise_vec(x): """ -def piecewise(x): - if x < -1: - return -1.0 - elif x > 1: - return 1.0 - return x -""" + def piecewise(x): + if x < -1: + return -1.0 + elif x > 1: + return 1.0 + return x + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x < -1, -1.0, xp.where(x > 1, 1.0, x)) diff --git a/tests/golden/cases/pow_op.py b/tests/golden/cases/pow_op.py index d7bde61..9872b9f 100644 --- a/tests/golden/cases/pow_op.py +++ b/tests/golden/cases/pow_op.py @@ -2,8 +2,8 @@ def pow_op_vec(x): """ -def pow_op(x): - return x**2 + 2.0**x -""" + def pow_op(x): + return x**2 + 2.0**x + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.pow(x, 2) + xp.pow(2.0, x) diff --git a/tests/golden/cases/psi.py b/tests/golden/cases/psi.py index 7269171..d4aa00c 100644 --- a/tests/golden/cases/psi.py +++ b/tests/golden/cases/psi.py @@ -2,10 +2,10 @@ def psi_vec(x): """ -def psi(x): - if x < 0: - return 0.0 - return x * math.exp(-x) -""" + def psi(x): + if x < 0: + return 0.0 + 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(xp.astype(xp.asarray(-x), xp.float64))) diff --git a/tests/golden/cases/relu.py b/tests/golden/cases/relu.py index 7fb574a..54e4361 100644 --- a/tests/golden/cases/relu.py +++ b/tests/golden/cases/relu.py @@ -2,10 +2,10 @@ def relu_vec(x): """ -def relu(x): - if x > 0: - return x - return 0.0 -""" + def relu(x): + if x > 0: + return x + return 0.0 + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, x, 0.0) diff --git a/tests/golden/cases/relu_with_else.py b/tests/golden/cases/relu_with_else.py index b3a8039..7eb914f 100644 --- a/tests/golden/cases/relu_with_else.py +++ b/tests/golden/cases/relu_with_else.py @@ -2,13 +2,13 @@ def relu_with_else_vec(x): """ -def relu_with_else(x): - if x > 0: # noqa: SIM108 - r = x - else: - r = 0.0 - return r -""" + def relu_with_else(x): + if x > 0: # noqa: SIM108 + r = x + else: + r = 0.0 + return r + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) r_1 = x r_2 = 0.0 diff --git a/tests/golden/cases/reserved_param.py b/tests/golden/cases/reserved_param.py index 2f48fa3..e1d1464 100644 --- a/tests/golden/cases/reserved_param.py +++ b/tests/golden/cases/reserved_param.py @@ -2,8 +2,8 @@ def reserved_param_vec(xp): """ -def reserved_param(xp): - return xp + 1 -""" + def reserved_param(xp): + return xp + 1 + """ xp_1 = array_namespace(*[a for a in (xp,) if hasattr(a, '__array_namespace__')]) return xp + 1 diff --git a/tests/golden/cases/ssa_rebind.py b/tests/golden/cases/ssa_rebind.py index e49bc50..e9d82fa 100644 --- a/tests/golden/cases/ssa_rebind.py +++ b/tests/golden/cases/ssa_rebind.py @@ -2,12 +2,12 @@ def ssa_rebind_vec(x): """ -def ssa_rebind(x): - y = x + 1 - y = y * 2 - y += 3 - return y -""" + def ssa_rebind(x): + y = x + 1 + y = y * 2 + y += 3 + return y + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) y = x + 1 y_1 = y * 2 diff --git a/tests/golden/cases/ternary.py b/tests/golden/cases/ternary.py index 4d1aa3f..8e0f2a6 100644 --- a/tests/golden/cases/ternary.py +++ b/tests/golden/cases/ternary.py @@ -2,8 +2,8 @@ def ternary_vec(x): """ -def ternary(x): - return x * 2 if x > 0 else x / 2 -""" + def ternary(x): + return x * 2 if x > 0 else x / 2 + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return xp.where(x > 0, x * 2, x / 2) diff --git a/tests/golden/cases/unary_ops.py b/tests/golden/cases/unary_ops.py index 7dba06a..a69b032 100644 --- a/tests/golden/cases/unary_ops.py +++ b/tests/golden/cases/unary_ops.py @@ -2,8 +2,8 @@ def unary_ops_vec(x): """ -def unary_ops(x): - return -x + +x - ~x -""" + def unary_ops(x): + return -x + +x - ~x + """ xp = array_namespace(*[a for a in (x,) if hasattr(a, '__array_namespace__')]) return -x + +x - ~x diff --git a/tests/test_inspectability.py b/tests/test_inspectability.py index e605f5e..5ca36ca 100644 --- a/tests/test_inspectability.py +++ b/tests/test_inspectability.py @@ -2,6 +2,7 @@ from __future__ import annotations +import ast import inspect import linecache import traceback @@ -86,6 +87,17 @@ def test_name_and_docstring() -> None: assert "if x < 0:" in (vec.__doc__ or "") +def test_docstring_cleandoc_exact() -> None: + # the docstring indents the original source to body level; cleandoc + # (ast.get_docstring / inspect.getdoc) recovers the original verbatim + vec = vectorize(relu_fn) + fn_def = ast.parse(vec.source).body[-1] + assert isinstance(fn_def, ast.FunctionDef) + original = inspect.getsource(relu_fn).strip() + assert ast.get_docstring(fn_def) == original + assert inspect.getdoc(vec) == original + + def test_idempotent_revectorization() -> None: vec1 = vectorize(relu_fn) vec2 = vectorize(vec1) diff --git a/tests/unit/test_codegen.py b/tests/unit/test_codegen.py index 8611754..16dd6f1 100644 --- a/tests/unit/test_codegen.py +++ b/tests/unit/test_codegen.py @@ -51,7 +51,7 @@ def test_module_shape() -> None: src = gen(Program(("x",), (), Ref("x"))) assert src.startswith("from array_api_compat import array_namespace") assert "def f_vec(x):" in src - assert '"""\ndef f(x):\n return x\n"""' in src + assert '"""\n def f(x):\n return x\n """' in src assert "xp = array_namespace(" in src assert "hasattr(a, '__array_namespace__')" in src