def testScanTypeErrors(self):
    'Test typing error messages for scan.'
    a = jnp.arange(5)
    with self.assertRaisesRegex(TypeError, re.escape('scan body output must be a pair, got ShapedArray(float32[]).')):
        lax.scan((lambda c, x: np.float32(0.0)), 0, a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure((0, 0, 0))} and {tree_util.tree_structure((1, (2, 3)))}')):
        lax.scan((lambda c, x: ((0, 0, 0), x)), (1, (2, 3)), a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure(a)} and {tree_util.tree_structure(None)}.')):
        lax.scan((lambda c, x: (0, x)), None, a)
    with self.assertRaisesWithLiteralMatch(TypeError, 'scan carry output and input must have identical types, got\nShapedArray(int32[])\nand\nShapedArray(float32[]).'):
        lax.scan((lambda c, x: (np.int32(0), x)), np.float32(1.0), a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure(a)} and {tree_util.tree_structure((1, 2))}.')):
        lax.scan((lambda c, x: (0, x)), (1, 2), a)
Test typing error messages for scan.
def testScanTypeErrors(self):
    
    a = jnp.arange(5)
    with self.assertRaisesRegex(TypeError, re.escape('scan body output must be a pair, got ShapedArray(float32[]).')):
        lax.scan((lambda c, x: np.float32(0.0)), 0, a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure((0, 0, 0))} and {tree_util.tree_structure((1, (2, 3)))}')):
        lax.scan((lambda c, x: ((0, 0, 0), x)), (1, (2, 3)), a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure(a)} and {tree_util.tree_structure(None)}.')):
        lax.scan((lambda c, x: (0, x)), None, a)
    with self.assertRaisesWithLiteralMatch(TypeError, 'scan carry output and input must have identical types, got\nShapedArray(int32[])\nand\nShapedArray(float32[]).'):
        lax.scan((lambda c, x: (np.int32(0), x)), np.float32(1.0), a)
    with self.assertRaisesRegex(TypeError, re.escape(f'scan carry output and input must have same type structure, got {tree_util.tree_structure(a)} and {tree_util.tree_structure((1, 2))}.')):
        lax.scan((lambda c, x: (0, x)), (1, 2), a)