def testSwitchErrors(self):
    'Test typing error messages for switch.'
    with self.assertRaisesRegex(TypeError, re.escape('Index type must be an integer, got <function')):
        lax.switch((lambda x: True), [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape('Index type must be an integer, got foo.')):
        lax.switch('foo', [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape('Branch index must be scalar, got (1.0, 1.0) of shape (2,).')):
        lax.switch((1.0, 1.0), [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(ValueError, re.escape('Empty branch sequence')):
        lax.switch(0, [], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'branch 0 and 1 outputs must have same type structure, got {tree_util.tree_structure(2.0)} and {tree_util.tree_structure((3.0, 3.0))}.')):
        lax.switch(1, [(lambda _: 2.0), (lambda _: (3.0, 3.0))], 1.0)
    with self.assertRaisesRegex(TypeError, textwrap.dedent('\n            branch 0 and 1 outputs must have identical types, got\n            ShapedArray\\(float32\\[1\\]\\)\n            and\n            ShapedArray\\(float32\\[\\].*\\).').strip()):
        lax.switch(1, [(lambda _: jnp.array([1.0], jnp.float32)), (lambda _: jnp.float32(1.0))], 1.0)
Test typing error messages for switch.
def testSwitchErrors(self):
    
    with self.assertRaisesRegex(TypeError, re.escape('Index type must be an integer, got <function')):
        lax.switch((lambda x: True), [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape('Index type must be an integer, got foo.')):
        lax.switch('foo', [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape('Branch index must be scalar, got (1.0, 1.0) of shape (2,).')):
        lax.switch((1.0, 1.0), [(lambda _: 2.0), (lambda _: 3.0)], 1.0)
    with self.assertRaisesRegex(ValueError, re.escape('Empty branch sequence')):
        lax.switch(0, [], 1.0)
    with self.assertRaisesRegex(TypeError, re.escape(f'branch 0 and 1 outputs must have same type structure, got {tree_util.tree_structure(2.0)} and {tree_util.tree_structure((3.0, 3.0))}.')):
        lax.switch(1, [(lambda _: 2.0), (lambda _: (3.0, 3.0))], 1.0)
    with self.assertRaisesRegex(TypeError, textwrap.dedent('\n            branch 0 and 1 outputs must have identical types, got\n            ShapedArray\\(float32\\[1\\]\\)\n            and\n            ShapedArray\\(float32\\[\\].*\\).').strip()):
        lax.switch(1, [(lambda _: jnp.array([1.0], jnp.float32)), (lambda _: jnp.float32(1.0))], 1.0)