def testScanLengthArg(self):

    def arange(n):
        return lax.scan((lambda c, _: ((c + 1), c)), 0, None, length=n)[1]
    ans = arange(10)
    expected = np.arange(10)
    self.assertAllClose(ans, expected, check_dtypes=False)
def testScanLengthArg(self):

    def arange(n):
        return lax.scan((lambda c, _: ((c + 1), c)), 0, None, length=n)[1]
    ans = arange(10)
    expected = np.arange(10)
    self.assertAllClose(ans, expected, check_dtypes=False)