We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
Element-wise Operations
Apply the same polynomial independently to every pair of corresponding values in two tensors.
Your task
Implement elementwise_poly(a, b) using the expanded expression:
a**2 + 2*a*b + b**2
The floating-point inputs have the same shape, dtype and device. Return that shape, dtype and device, leaving both inputs unchanged. No reduction over an axis or matrix multiplication is needed.
For a=[1, -2] and b=[3, 1], return [16, 1]:
- First position:
1**2 + 2*1*3 + 3**2 = 16. - Second position:
(-2)**2 + 2*(-2)*1 + 1**2 = 1.
The same rule applies to matrices and higher-dimensional tensors. Each output
uses only a and b at that position. In both frameworks, * multiplies values
at matching positions, while **2 squares each value. Python's ^ is not a
power operator.
Algebraically the polynomial is (a + b)**2, but this exercise asks for the
expanded form to practice elementwise addition, multiplication and powers.
Floating-point evaluation can give slightly different rounding for the two forms.
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Choose matching entries for arithmetic
Elementwise operations act on corresponding entries, after any broadcasting.
For a = [1, 2] and b = [3, 4], (a + b)^2 gives [16, 36]. A dot product
instead multiplies corresponding entries and sums them, giving 11.
Broadcasting can make different shapes compatible. Check that it pairs the intended entries; an operation can run successfully and still compute the wrong thing.
Select the largest values along the intended axis
Top-k selection keeps the largest k entries along an axis. The top two values
of [4, 1, 3], in descending order, are [4, 3]. Their original indices are
[0, 2]. Values and indices are different outputs, so return what the task asks
for.
Check three details: which axis to select along, whether the result must be
sorted, and how to resolve ties. A partial-selection algorithm can find the top
entries without sorting them. Test k = 1 and k equal to the axis length, and
use the task's stated behavior for invalid or empty inputs.
Preserve the requested number type
Mathematically equivalent expressions can give slightly different floating-point results. Integer arithmetic can also overflow before a later conversion to floating point. Follow the Problem tab's requirements for dtype, precision, and output shape.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
JAX reference solution
Sign in to track your attempts toward unlocking the reference solution.