@parameterized.named_parameters(({'testcase_name': f'_{shape}_axis={axis}', 'shape': shape, 'axis': axis} for shape in [[0], [1], [2], [3], [5], [10], [1000], [2, 3], [7, 5], [5, 6, 7]] for axis in range((- len(shape)), (len(shape) - 1))))
def testAssociativeScanUnstructured(self, shape, axis):
    data = (np.arange(np.prod(shape)).reshape(shape) + 7)
    expected = np.cumsum(data, axis=axis)
    result = lax.associative_scan(operator.add, data, axis=axis)
    self.assertAllClose(result, expected, check_dtypes=False)
@parameterized.named_parameters(({'testcase_name': f'_{shape}_axis={axis}', 'shape': shape, 'axis': axis} for shape in [[0], [1], [2], [3], [5], [10], [1000], [2, 3], [7, 5], [5, 6, 7]] for axis in range((- len(shape)), (len(shape) - 1))))
def testAssociativeScanUnstructured(self, shape, axis):
    data = (np.arange(np.prod(shape)).reshape(shape) + 7)
    expected = np.cumsum(data, axis=axis)
    result = lax.associative_scan(operator.add, data, axis=axis)
    self.assertAllClose(result, expected, check_dtypes=False)