def test_cond_typecheck_param(self):

    def new_jaxpr():
        jaxpr = api.make_jaxpr((lambda x: lax.switch(0, [jnp.sin, jnp.cos], x)))(1.0).jaxpr
        cond = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'cond')))
        return (jaxpr, cond)
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['branches'] = (4, 2)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid cond param branches of type tuple, tuple of ClosedJaxpr required: (4, 2)'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['linear'] = (4, 2)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid cond param linear of type tuple, tuple of bool required: (4, 2)'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['linear'] = 'multi\nline'
    self.assertRaisesRegex(core.JaxprTypeError, 'invalid cond param linear of type str, tuple of bool required:\\nmulti\\nline', (lambda : core.check_jaxpr(jaxpr)))
def test_cond_typecheck_param(self):

    def new_jaxpr():
        jaxpr = api.make_jaxpr((lambda x: lax.switch(0, [jnp.sin, jnp.cos], x)))(1.0).jaxpr
        cond = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'cond')))
        return (jaxpr, cond)
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['branches'] = (4, 2)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid cond param branches of type tuple, tuple of ClosedJaxpr required: (4, 2)'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['linear'] = (4, 2)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid cond param linear of type tuple, tuple of bool required: (4, 2)'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['linear'] = 'multi\nline'
    self.assertRaisesRegex(core.JaxprTypeError, 'invalid cond param linear of type str, tuple of bool required:\\nmulti\\nline', (lambda : core.check_jaxpr(jaxpr)))