def testAssociativeScanStructured3(self):
    pair = collections.namedtuple('pair', ('first', 'second'))
    data = pair(first=np.array([0.0, 1.0, 2.0]), second=np.array([0.0, 10.0, 20.0]))

    def fn(a, b):
        return pair(first=(a.first + b.first), second=(a.second + b.second))
    result = lax.associative_scan(fn, elems=data)
    self.assertAllClose(result.first, np.array([0.0, 1.0, 3.0]), check_dtypes=False)
    self.assertAllClose(result.second, np.array([0.0, 10.0, 30.0]), check_dtypes=False)
def testAssociativeScanStructured3(self):
    pair = collections.namedtuple('pair', ('first', 'second'))
    data = pair(first=np.array([0.0, 1.0, 2.0]), second=np.array([0.0, 10.0, 20.0]))

    def fn(a, b):
        return pair(first=(a.first + b.first), second=(a.second + b.second))
    result = lax.associative_scan(fn, elems=data)
    self.assertAllClose(result.first, np.array([0.0, 1.0, 3.0]), check_dtypes=False)
    self.assertAllClose(result.second, np.array([0.0, 10.0, 30.0]), check_dtypes=False)