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

    def f(c, a):
        assert (a.shape == (3,))
        assert (c.shape == (4,))
        b = jnp.cos(((jnp.sum(jnp.sin(a)) + jnp.sum(jnp.cos(c))) + jnp.sum(jnp.tan(d))))
        c = jnp.sin((c * b))
        assert (b.shape == ())
        return (c, b)
    xs = jnp.ones((5, 3))
    c = jnp.ones(4)
    scan = (lambda c, xs: lax.scan(f, c, xs))
    scan_unrolled = (lambda c, xs: lax.scan(f, c, xs, unroll=2))
    self.assertEqual(len(str(api.make_jaxpr(scan)(c, xs))), len(str(api.make_jaxpr(scan_unrolled)(c, xs))))
    self.assertLess(len(str(api.xla_computation(scan)(c, xs).as_hlo_text())), len(str(api.xla_computation(scan_unrolled)(c, xs).as_hlo_text())))
def test_scan_unroll(self):
    d = jnp.ones(2)

    def f(c, a):
        assert (a.shape == (3,))
        assert (c.shape == (4,))
        b = jnp.cos(((jnp.sum(jnp.sin(a)) + jnp.sum(jnp.cos(c))) + jnp.sum(jnp.tan(d))))
        c = jnp.sin((c * b))
        assert (b.shape == ())
        return (c, b)
    xs = jnp.ones((5, 3))
    c = jnp.ones(4)
    scan = (lambda c, xs: lax.scan(f, c, xs))
    scan_unrolled = (lambda c, xs: lax.scan(f, c, xs, unroll=2))
    self.assertEqual(len(str(api.make_jaxpr(scan)(c, xs))), len(str(api.make_jaxpr(scan_unrolled)(c, xs))))
    self.assertLess(len(str(api.xla_computation(scan)(c, xs).as_hlo_text())), len(str(api.xla_computation(scan_unrolled)(c, xs).as_hlo_text())))