def test_amax():
inp = np.zeros((INT_OVERFLOW, 2))
inp[-1, -1] = 1
inp.attach_grad()
with mx.autograd.record():
out = np.amax(inp)
out.backward()
assert out == 1.0
assert inp.grad.shape == (INT_OVERFLOW, 2)
assert inp.grad[0, 0] == 0 and inp.grad[-1, -1] == 1
def test_amax():
A = np.zeros((INT_OVERFLOW, 2))
A[100][1] = 1
A.attach_grad()
with mx.autograd.record():
B = np.amax(A)
print(B)
assert B == 1.0
B.backward()
assert A.grad.shape == (INT_OVERFLOW, 2)
assert A.grad[0][0] == 0