def test_setitem(self):
# Single integer index.
a = array_creation.array([1., 2., 3.])
b = array_creation.array(5.)
c = array_creation.array(10.)
tensors = [arr.data for arr in [a, b, c]]
with tf.GradientTape() as g:
g.watch(tensors)
a[1] = b + c
loss = array_methods.sum(a)
gradients = g.gradient(loss.data, tensors)
self.assertSequenceEqual(
array_creation.array(gradients[0]).tolist(), [1., 0., 1.])
self.assertEqual(array_creation.array(gradients[1]).tolist(), 1.)
self.assertEqual(array_creation.array(gradients[2]).tolist(), 1.)
# Tuple index.
a = array_creation.array([[[1., 2.], [3., 4.]],
[[5., 6.], [7., 8.]]]) # 2x2x2 array.
b = array_creation.array([10., 11.])
tensors = [arr.data for arr in [a, b]]
with tf.GradientTape() as g:
g.watch(tensors)
a[(1, 0)] = b
loss = array_methods.sum(a)
gradients = g.gradient(loss.data, tensors)
self.assertSequenceEqual(
array_creation.array(gradients[0]).tolist(),
[[[1., 1.], [1., 1.]], [[0., 0.], [1., 1.]]])
self.assertEqual(array_creation.array(gradients[1]).tolist(), [1., 1.])
def nanmean(a, axis=None, dtype=None, keepdims=None):
a = array_creation.asarray(a)
if np.issubdtype(a.dtype, np.bool_) or np.issubdtype(a.dtype, np.integer):
return array_methods.mean(a, axis=axis, dtype=dtype, keepdims=keepdims)
nan_mask = logical_not(isnan(a))
normalizer = array_methods.sum(nan_mask, axis=axis, dtype=np.int64,
keepdims=keepdims)
return nansum(a, axis=axis, dtype=dtype, keepdims=keepdims) / normalizer
def f(a, b): return array_methods.sum(math.sqrt(math.exp(a)) + b)
