def testCondTypeErrors(self):
    'Test typing error messages for  cond.'
    with self.assertRaisesRegex(TypeError, re.escape('Pred type must be either boolean or number, got <function')):
        lax.cond((lambda x: True), (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape("Pred must be a scalar, got foo of type <class 'str'>")):
        lax.cond('foo', (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape("Pred must be a scalar, got (1.0, 1.0) of type <class 'tuple'>")):
        lax.cond((1.0, 1.0), (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'true_fun and false_fun output must have same type structure, got {tree_util.tree_structure(2.0)} and {tree_util.tree_structure((3.0, 3.0))}.')):
        lax.cond(True, (lambda top: 2.0), (lambda fop: (3.0, 3.0)), 1.0)
    with self.assertRaisesRegex(TypeError, textwrap.dedent('\n            true_fun and false_fun output must have identical types, got\n            ShapedArray\\(float32\\[1\\]\\)\n            and\n            ShapedArray\\(float32\\[\\].*\\).').strip()):
        lax.cond(True, (lambda top: jnp.array([1.0], jnp.float32)), (lambda fop: jnp.float32(1.0)), 1.0)
Test typing error messages for  cond.
def testCondTypeErrors(self):
    
    with self.assertRaisesRegex(TypeError, re.escape('Pred type must be either boolean or number, got <function')):
        lax.cond((lambda x: True), (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape("Pred must be a scalar, got foo of type <class 'str'>")):
        lax.cond('foo', (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape("Pred must be a scalar, got (1.0, 1.0) of type <class 'tuple'>")):
        lax.cond((1.0, 1.0), (lambda top: 2.0), (lambda fop: 3.0), 1.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'true_fun and false_fun output must have same type structure, got {tree_util.tree_structure(2.0)} and {tree_util.tree_structure((3.0, 3.0))}.')):
        lax.cond(True, (lambda top: 2.0), (lambda fop: (3.0, 3.0)), 1.0)
    with self.assertRaisesRegex(TypeError, textwrap.dedent('\n            true_fun and false_fun output must have identical types, got\n            ShapedArray\\(float32\\[1\\]\\)\n            and\n            ShapedArray\\(float32\\[\\].*\\).').strip()):
        lax.cond(True, (lambda top: jnp.array([1.0], jnp.float32)), (lambda fop: jnp.float32(1.0)), 1.0)