compute a * x + y
a (Any) – scalar
x (Any) – pytrees with the same treedef
y (Any) – pytrees with the same treedef
The sum of the respective leaves of the two pytrees x and y
where the leaves of x are first scaled with a.