Multivariate lambda function in Python that scales with number of input variables received

Question

The following toy function ordinarily takes two input variables:

f = lambda u1, u2 : (u1*u2)*(u1**2+u2**2)

but can scale beyond the bivariate case to higher dimensions:

if dim == 2:
    f = lambda u1, u2 : (u1*u2)*(u1**2+u2**2)
if dim == 3:
    f = lambda u1, u2, u3 : (u1*u2*u3)*(u1**2+u2**2+u3**2)
if dim == 4:
    f = lambda u1, u2, u3, u4 : (u1*u2*u3*u4)*(u1**2+u2**2+u3**2+u4**2)

How can the lambda function be written so that it can expand itself in the call lambda u1, u2, u3, u4, ... as well as the function body itself, based on number of inputs sent to it, similar to how a defined function can be declared as def f(*args) where *args is an arbitrary number of input arguments?

Answer 1

The lambda syntax supports the same parameter list syntax as the def syntax, including variadic positional and keyword argument.

f = lambda *us: math.prod(us) * sum(u**2 for u in us)

If the *us are not invariant when multiplied by 1 or added to 0, the * and + operations can be applied across elements via reduce:

from functools import reduce
import operator

f = lambda *us: reduce(operator.mul, us) * reduce(operator.add, (u**2 for u in us))