at.square(x) and at.pow(x, 2) don't etuplize to the same expresssion

[rlouf]

import aesara.tensor as at
from etuples import etuplize

a = at.scalar('a')

etuplize(at.square(a))
# e(e(<class 'aesara.tensor.elemwise.Elemwise'>, sqr, <frozendict {}>), a)

etuplize(at.pow(a, 2))
# e(e(<class 'aesara.tensor.elemwise.Elemwise'>, pow, <frozendict {}>), a, TensorConstant{2})

Which means that a**2 will fail to unify with etuplize(at.square(a)):

print(etuplize(a ** 2))
# e(e(<class 'aesara.tensor.elemwise.Elemwise'>, pow, <frozendict {}>), a, TensorConstant{2})

I guess this question is more about why at.square is not an alias to at.pow(..., 2).

[ricardoV94]

I guess this question is more about why at.square is not an alias to at.pow(..., 2).

Probably optimization. Wouldn't be surprised if squaring was faster than power(x, 2)

[rlouf]

So that could eventually become a rewrite? And is that universally true for every backend who is that distinction tied to the C backend?

[ricardoV94]

We already have some rewrites:

aesara/aesara/tensor/rewriting/math.py

Line 1892 in ec82b9f

def local_pow_specialize(fgraph, node):

I guess you mean we might be missing an intermediate canonicalization form that is the same for either graph

[rlouf]

I guess you mean we might be missing an intermediate canonicalization form that is the same for either graph

Yes.

Then we can have at.square(x) be an alias for at.pow(x, 2), which is a question of consistency of the representation in the IR. And then let the canonicalisation handle the computation cost concerns.

[rlouf]

Note that at.reciprocal(x) should be an alias for at.pow(x, -1) and at.sqrt(x) an alias for at.pow(x, 0.5) for the same consistency reasons.

Consistency is going to prevent special-casing the code in downstream libraries, for instance in AePPL.