pytaco.square

pytaco.square(e1)

Return the element-wise square value of the index expression.

This must be assigned to a tensor for the computation to be performed.

Parameters

e1: index_expression

Input index expression

Returns

sq_exp: index_expression

An index expression representing the element wise square of the input expression.

Examples

>>> import pytaco as pt
>>> t = pt.as_tensor([-2, 2, 1])
>>> i = pt.index_var()
>>> sq_expr = pt.square(t[i])

We can then assign this description to a tensor to actually perform the computation.

The code below tells taco to compute the square of each value, sum over all those values and store it in the tensor res_t. Since i appears on the right hand side of the expression but not on the left, taco will take the sum of the values produced.

>>> res_t = pt.tensor()
>>> res_t[None] = sq_expr
>>> res_t.to_array()
array(9., dtype=float32)