Element-wise Operations

Given two tensors a and b of the same shape, compute:

$$\text{result} = a^2 + 2ab + b^2$$

This is equivalent to (a + b)^2 but implement it using individual element-wise operations.

Input: Tensors a and b of the same shape.

Output: A tensor of the same shape equal to a**2 + 2*a*b + b**2.

API Reference:

• PyTorch: torch.pow, torch.mul, torch.add, or operators +, *, **
• JAX: jnp.power, jnp.multiply, jnp.add, or operators