def test_scan_typecheck_param(self):
    d = jnp.ones(2)

    def f(c, a):
        b = jnp.cos(((jnp.sum(a) + jnp.sum(c)) + jnp.sum(d)))
        c = jnp.sin((c * b))
        return (c, b)
    xs = jnp.ones((5, 3))
    c = jnp.ones(4)
    scan_fun = (lambda c, xs: lax.scan(f, c, xs))

    def new_jaxpr():
        jaxpr = api.make_jaxpr(scan_fun)(c, xs).jaxpr
        scan = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'scan')))
        return (jaxpr, scan)
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['reverse'] = 4
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid scan param reverse of type int, bool required: 4'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['num_consts'] = (- 3)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid scan param num_consts of type int, non-negative int required: -3'), (lambda : core.check_jaxpr(jaxpr)))
def test_scan_typecheck_param(self):
    d = jnp.ones(2)

    def f(c, a):
        b = jnp.cos(((jnp.sum(a) + jnp.sum(c)) + jnp.sum(d)))
        c = jnp.sin((c * b))
        return (c, b)
    xs = jnp.ones((5, 3))
    c = jnp.ones(4)
    scan_fun = (lambda c, xs: lax.scan(f, c, xs))

    def new_jaxpr():
        jaxpr = api.make_jaxpr(scan_fun)(c, xs).jaxpr
        scan = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'scan')))
        return (jaxpr, scan)
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['reverse'] = 4
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid scan param reverse of type int, bool required: 4'), (lambda : core.check_jaxpr(jaxpr)))
    (jaxpr, eqn) = new_jaxpr()
    eqn.params['num_consts'] = (- 3)
    self.assertRaisesRegex(core.JaxprTypeError, re.escape('invalid scan param num_consts of type int, non-negative int required: -3'), (lambda : core.check_jaxpr(jaxpr)))