def testWhileTypeErrors(self):
    'Test typing error messages for while.'
    tuple_treedef = tree_util.tree_structure((1.0, 1.0))
    leaf_treedef = tree_util.tree_structure(0.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'cond_fun must return a boolean scalar, but got pytree {tuple_treedef}.')):
        lax.while_loop((lambda c: (1.0, 1.0)), (lambda c: c), 0.0)
    with self.assertRaisesRegex(TypeError, re.escape('cond_fun must return a boolean scalar, but got output type(s) [ShapedArray(float32[])].')):
        lax.while_loop((lambda c: np.float32(1.0)), (lambda c: c), np.float32(0.0))
    with self.assertRaisesRegex(TypeError, re.escape(f'body_fun output and input must have same type structure, got {tuple_treedef} and {leaf_treedef}.')):
        lax.while_loop((lambda c: True), (lambda c: (1.0, 1.0)), 0.0)
    with self.assertRaisesWithLiteralMatch(TypeError, 'body_fun output and input must have identical types, got\nShapedArray(bool[], weak_type=True)\nand\nShapedArray(float32[]).'):
        lax.while_loop((lambda c: True), (lambda c: True), np.float32(0.0))
Test typing error messages for while.
def testWhileTypeErrors(self):
    
    tuple_treedef = tree_util.tree_structure((1.0, 1.0))
    leaf_treedef = tree_util.tree_structure(0.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'cond_fun must return a boolean scalar, but got pytree {tuple_treedef}.')):
        lax.while_loop((lambda c: (1.0, 1.0)), (lambda c: c), 0.0)
    with self.assertRaisesRegex(TypeError, re.escape('cond_fun must return a boolean scalar, but got output type(s) [ShapedArray(float32[])].')):
        lax.while_loop((lambda c: np.float32(1.0)), (lambda c: c), np.float32(0.0))
    with self.assertRaisesRegex(TypeError, re.escape(f'body_fun output and input must have same type structure, got {tuple_treedef} and {leaf_treedef}.')):
        lax.while_loop((lambda c: True), (lambda c: (1.0, 1.0)), 0.0)
    with self.assertRaisesWithLiteralMatch(TypeError, 'body_fun output and input must have identical types, got\nShapedArray(bool[], weak_type=True)\nand\nShapedArray(float32[]).'):
        lax.while_loop((lambda c: True), (lambda c: True), np.float32(0.0))