Exemplo n.º 1
0
def test51_scatter_reduce_fwd_eager(m):
    with EagerMode():
        for i in range(3):
            idx1 = ek.arange(m.UInt, 5)
            idx2 = ek.arange(m.UInt, 4) + 3

            x = ek.linspace(m.Float, 0, 1, 5)
            y = ek.linspace(m.Float, 1, 2, 4)
            buf = ek.zero(m.Float, 10)

            if i % 2 == 0:
                ek.enable_grad(buf)
                ek.set_grad(buf, 1)
            if i // 2 == 0:
                ek.enable_grad(x, y)
                ek.set_grad(x, 1)
                ek.set_grad(y, 1)

            x.label = "x"
            y.label = "y"
            buf.label = "buf"

            buf2 = m.Float(buf)
            ek.scatter_reduce(ek.ReduceOp.Add, buf2, x, idx1)
            ek.scatter_reduce(ek.ReduceOp.Add, buf2, y, idx2)

            s = ek.dot_async(buf2, buf2)

            # Verified against Mathematica
            assert ek.allclose(ek.detach(s), 15.5972)
            assert ek.allclose(ek.grad(s), (25.1667 if i // 2 == 0 else 0) +
                               (17 if i % 2 == 0 else 0))
Exemplo n.º 2
0
def test16_custom(cname):
    t = get_class(cname)

    v1 = ek.zero(t, 100)
    v2 = ek.empty(t, 100)
    assert len(v1.state) == 100
    assert len(v2.inc) == 100

    v2.state = v1.state
    v1.state = ek.arange(type(v1.state), 100)
    v3 = ek.select(v1.state < 10, v1, v2)
    assert v3.state[3] == 3
    assert v3.state[11] == 0

    assert ek.width(v3) == 100
    v4 = ek.zero(t, 1)
    ek.schedule(v4)
    ek.resize(v4, 200)
    assert ek.width(v4) == 200

    assert ek.width(v3) == 100
    v4 = ek.zero(t, 1)
    ek.resize(v4, 200)
    assert ek.width(v4) == 200

    index = ek.arange(type(v1.state), 100)
    ek.scatter(v4, v1, index)
    v5 = ek.gather(t, v4, index)
    ek.eval(v5)
    assert v5.state == v1.state and v5.inc == v1.inc
Exemplo n.º 3
0
def test_roundtrip_dlpack_all(package):
    Float, Array3f = package.Float, package.Array3f
    prepare(package)
    a1 = Array3f(
        ek.arange(Float, 10),
        ek.arange(Float, 10) + 100,
        ek.arange(Float, 10) + 1000,
    )
    a2 = a1.dlpack()
    a3 = Array3f(a2)
    assert a1 == a3
    assert a1.x == Float(a1.x.dlpack())
Exemplo n.º 4
0
def test_roundtrip_numpy_all(package):
    pytest.importorskip("numpy")
    Float, Array3f = package.Float, package.Array3f
    prepare(package)
    a1 = Array3f(
        ek.arange(Float, 10),
        ek.arange(Float, 10) + 100,
        ek.arange(Float, 10) + 1000,
    )
    a2 = a1.numpy()
    a3 = Array3f(a2)
    assert a1 == a3
    assert a1.x == Float(a1.x.numpy())
Exemplo n.º 5
0
def test_roundtrip_pytorch_llvm():
    pytest.importorskip("torch")
    prepare(enoki.llvm.ad)
    from enoki.llvm.ad import Float, Array3f
    a1 = Array3f(
        ek.arange(Float, 10),
        ek.arange(Float, 10) + 100,
        ek.arange(Float, 10) + 1000,
    )
    a2 = a1.torch()
    a3 = Array3f(a2)
    assert a1 == a3
    assert a1.x == Float(a1.x.torch())
Exemplo n.º 6
0
def test_roundtrip_pytorch_jax():
    pytest.importorskip("jax")
    prepare(enoki.cuda.ad)
    from enoki.cuda.ad import Float, Array3f
    a1 = Array3f(
        ek.arange(Float, 10),
        ek.arange(Float, 10) + 100,
        ek.arange(Float, 10) + 1000,
    )
    a2 = a1.jax()
    a3 = Array3f(a2)
    assert a1 == a3
    assert a1.x == Float(a1.x.jax())
Exemplo n.º 7
0
def test20_scatter_add_rev(m):
    for i in range(3):
        idx1 = ek.arange(m.UInt, 5)
        idx2 = ek.arange(m.UInt, 4) + 3

        x = ek.linspace(m.Float, 0, 1, 5)
        y = ek.linspace(m.Float, 1, 2, 4)
        buf = ek.zero(m.Float, 10)

        if i % 2 == 0:
            ek.enable_grad(buf)
        if i // 2 == 0:
            ek.enable_grad(x, y)

        x.label = "x"
        y.label = "y"
        buf.label = "buf"

        buf2 = m.Float(buf)
        ek.scatter_add(buf2, x, idx1)
        ek.scatter_add(buf2, y, idx2)

        ref_buf = m.Float(0.0000, 0.2500, 0.5000, 1.7500, 2.3333, 1.6667,
                          2.0000, 0.0000, 0.0000, 0.0000)

        assert ek.allclose(ref_buf, buf2, atol=1e-4)
        assert ek.allclose(ref_buf, buf, atol=1e-4)

        s = ek.dot_async(buf2, buf2)

        print(ek.graphviz_str(s))

        ek.backward(s)

        ref_x = m.Float(0.0000, 0.5000, 1.0000, 3.5000, 4.6667)
        ref_y = m.Float(3.5000, 4.6667, 3.3333, 4.0000)

        if i // 2 == 0:
            assert ek.allclose(ek.grad(y), ek.detach(ref_y), atol=1e-4)
            assert ek.allclose(ek.grad(x), ek.detach(ref_x), atol=1e-4)
        else:
            assert ek.grad(x) == 0
            assert ek.grad(y) == 0

        if i % 2 == 0:
            assert ek.allclose(ek.grad(buf), ek.detach(ref_buf) * 2, atol=1e-4)
        else:
            assert ek.grad(buf) == 0
Exemplo n.º 8
0
def test22_scatter_fwd(m):
    x = m.Float(4.0)
    ek.enable_grad(x)

    values = x * x * ek.linspace(m.Float, 1, 4, 4)
    idx = 2 * ek.arange(m.UInt32, 4)

    buf = ek.zero(m.Float, 10)
    ek.scatter(buf, values, idx)

    assert ek.grad_enabled(buf)

    ref = [16.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0]
    assert ek.allclose(buf, ref)

    ek.forward(x, retain_graph=True)
    grad = ek.grad(buf)

    ref_grad = [8.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0]
    assert ek.allclose(grad, ref_grad)

    # Overwrite first value with non-diff value, resulting gradient entry should be 0
    y = m.Float(3)
    idx = m.UInt32(0)
    ek.scatter(buf, y, idx)

    ref = [3.0, 0.0, 32.0, 0.0, 48.0, 0.0, 64.0, 0.0, 0.0, 0.0]
    assert ek.allclose(buf, ref)

    ek.forward(x)
    grad = ek.grad(buf)

    ref_grad = [0.0, 0.0, 16.0, 0.0, 24.0, 0.0, 32.0, 0.0, 0.0, 0.0]
    assert ek.allclose(grad, ref_grad)
Exemplo n.º 9
0
        def eval(self, pos, vel):
            pos, vel = m.Array2f(pos), m.Array2f(vel)

            # Run for 100 iterations
            it, max_it = m.UInt32(0), 100

            # Allocate scratch space
            n = max(ek.width(pos), ek.width(vel))
            self.temp_pos = ek.empty(m.Array2f, n * max_it)
            self.temp_vel = ek.empty(m.Array2f, n * max_it)

            loop = m.Loop(pos, vel, it)
            while loop.cond(it < max_it):
                # Store current loop variables
                index = it * n + ek.arange(m.UInt32, n)
                ek.scatter(self.temp_pos, pos, index)
                ek.scatter(self.temp_vel, vel, index)

                # Update loop variables
                pos_out, vel_out = self.timestep(pos, vel)
                pos.assign(pos_out)
                vel.assign(vel_out)

                it += 1

            # Ensure output and temp. arrays are evaluated at this point
            ek.eval(pos, vel)

            return pos, vel
Exemplo n.º 10
0
def test02_multiple_values(pkg, variant):
    p = get_class(pkg)

    i = ek.arange(p.Int, 0, 10)
    v = ek.zero(p.Array3f, 10)

    if variant == 1:
        v.y = p.Float(0)

    loop = p.Loop(i, v)
    while loop.cond(i < 5):
        i.assign(i + 1)
        f = p.Float(i)
        v.x += f
        v.y += 2 * f
        v.z += 4 * f

    if variant == 0:
        ek.eval(i, v)
    else:
        ek.eval(i)
        ek.eval(v.x)
        ek.eval(v.y)
        ek.eval(v.z)

    assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
    assert v.y == p.Int(30, 28, 24, 18, 10, 0, 0, 0, 0, 0)
Exemplo n.º 11
0
        def backward(self):
            grad_pos, grad_vel = self.grad_out()

            # Run for 100 iterations
            it = m.UInt32(100)

            loop = m.Loop(it, grad_pos, grad_vel)
            n = ek.width(grad_pos)
            while loop.cond(it > 0):
                # Retrieve loop variables, reverse chronological order
                it -= 1
                index = it * n + ek.arange(m.UInt32, n)
                pos = ek.gather(m.Array2f, self.temp_pos, index)
                vel = ek.gather(m.Array2f, self.temp_vel, index)

                # Differentiate loop body in reverse mode
                ek.enable_grad(pos, vel)
                pos_out, vel_out = self.timestep(pos, vel)
                ek.set_grad(pos_out, grad_pos)
                ek.set_grad(vel_out, grad_vel)
                ek.enqueue(pos_out, vel_out)
                ek.traverse(m.Float, reverse=True)

                # Update loop variables
                grad_pos.assign(ek.grad(pos))
                grad_vel.assign(ek.grad(vel))

            self.set_grad_in('pos', grad_pos)
            self.set_grad_in('vel', grad_vel)
Exemplo n.º 12
0
def test01_ctr(pkg):
    p = get_class(pkg)

    i = ek.arange(p.Int, 0, 10)

    loop = p.Loop(i)
    while loop.cond(i < 5):
        i += 1

    assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
Exemplo n.º 13
0
def test14_roundtrip(package, dest):
    pytest.importorskip(dest)
    Float, Array3f, Array4f = package.Float, package.Array3f, package.Array4f
    Matrix3f, Matrix4f, Matrix44f = package.Matrix3f, package.Matrix4f, package.Matrix44f
    prepare(package)

    def to_dest(a):
        if dest == 'numpy':
            return a.numpy()
        if dest == 'torch':
            return a.torch()
        if dest == 'jax':
            return a.jax()

    v = Array3f(
        (1.0 + ek.arange(Float, 5)),
        (1.0 + ek.arange(Float, 5)) * 2,
        (1.0 + ek.arange(Float, 5)) * 3,
    )
    assert(v == Array3f(to_dest(v)))

    m = Matrix3f([
        Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]),
        Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]) * 10,
        Array3f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]) * 100
    ])
    assert(m ==Matrix3f(to_dest(m)))

    m = Matrix4f([
        Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]),
        Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 10,
        Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 100,
        Array4f([[1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]]) * 1000
    ])
    assert(m == Matrix4f(to_dest(m)))

    m = Matrix44f([
        [1.0, 2.0, 3.0, 4.0],
        [1.0, 2.0, 3.0, 4.0],
        [1.0, 2.0, 3.0, 4.0],
        [1.0, 2.0, 3.0, 4.0]
    ])
    assert(m == Matrix44f(to_dest(m)))
Exemplo n.º 14
0
def test02_radical_inverse_vectorized(variant_scalar_rgb):
    from mitsuba.core import RadicalInverse

    v = RadicalInverse()
    for index, prime in enumerate(gen_primes()):
        if index >= 1024:
            break
        result = v.eval(index, ek.arange(10, dtype=ek.uint64))
        for i in range(len(result)):
            assert ek.abs(r_inv(prime, i) - result[i]) < 1e-7
Exemplo n.º 15
0
def test01_ctr(pkg):
    p = get_class(pkg)

    i = ek.arange(p.Int, 0, 10)

    loop = p.Loop("MyLoop", lambda: i)
    while loop(i < 5):
        i = i + 1

    assert i == p.Int(5, 5, 5, 5, 5, 5, 6, 7, 8, 9)
Exemplo n.º 16
0
def test09_loop_simplification(pkg):
    # Test a regression where most loop variables are optimized away
    p = get_class(pkg)
    res = p.Float(0.0)
    active = p.Bool(True)
    depth = p.UInt32(1000)

    loop = p.Loop("TestLoop")
    loop.put(lambda: (active, depth, res))
    loop.init()
    while loop(active):
        res += ek.arange(p.Float, 10)
        depth -= 1
        active &= (depth > 0)
    del loop
    del active
    del depth
    gc.collect()
    gc.collect()
    assert res == ek.arange(p.Float, 10) * 1000
Exemplo n.º 17
0
def test07_gather_ravel_unravel(pkg):
    get_class(pkg.__name__)
    str_1 = '[[0.0, 1.0, 2.0],\n [3.0, 4.0, 5.0],\n [6.0, 7.0, 8.0],\n' \
        ' [9.0, 10.0, 11.0]]'
    str_2 = '[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]'
    a = ek.arange(pkg.Float, 100)
    b = ek.gather(pkg.Array3f, a, pkg.UInt(0, 1, 2, 3))
    assert repr(b) == str_1
    c = ek.ravel(b)
    assert repr(c) == str_2
    d = ek.unravel(pkg.Array3f, c)
    assert repr(d) == str_1
def test03_sample_eval_pdf(variant_scalar_rgb, interaction):
    from mitsuba.core.math import InvPi
    from mitsuba.core.warp import square_to_uniform_sphere
    from mitsuba.render import BSDFContext
    from mitsuba.core.xml import load_string

    bsdf = load_string("""<bsdf version="2.0.0" type="twosided">
        <bsdf type="diffuse">
            <rgb name="reflectance" value="0.1, 0.1, 0.1"/>
        </bsdf>
        <bsdf type="diffuse">
            <rgb name="reflectance" value="0.9, 0.9, 0.9"/>
        </bsdf>
    </bsdf>""")

    n = 5
    ctx = BSDFContext()
    for u in ek.arange(UInt32, n):
        for v in ek.arange(UInt32, n):
            interaction.wi = square_to_uniform_sphere([u / float(n-1),
                                                       v / float(n-1)])
            up = ek.dot(interaction.wi, [0, 0, 1]) > 0

            for x in ek.arange(UInt32, n):
                for y in ek.arange(UInt32, n):
                    sample = [x / float(n-1), y / float(n-1)]
                    (bs, s_value) = bsdf.sample(ctx, interaction, 0.5, sample)

                    if ek.any(s_value > 0):
                        # Multiply by square_to_cosine_hemisphere_theta
                        s_value *= bs.wo[2] * InvPi
                        if not up:
                            s_value *= -1

                        e_value = bsdf.eval(ctx, interaction, bs.wo)
                        p_pdf   = bsdf.pdf(ctx, interaction, bs.wo)

                        assert ek.allclose(s_value, e_value, atol=1e-2)
                        assert ek.allclose(bs.pdf, p_pdf)
                        assert not ek.any(ek.isnan(e_value) | ek.isnan(s_value))
Exemplo n.º 19
0
def test21_scatter_add_fwd(m):
    for i in range(3):
        idx1 = ek.arange(m.UInt, 5)
        idx2 = ek.arange(m.UInt, 4) + 3

        x = ek.linspace(m.Float, 0, 1, 5)
        y = ek.linspace(m.Float, 1, 2, 4)
        buf = ek.zero(m.Float, 10)

        if i % 2 == 0:
            ek.enable_grad(buf)
            ek.set_grad(buf, 1)
        if i // 2 == 0:
            ek.enable_grad(x, y)
            ek.set_grad(x, 1)
            ek.set_grad(y, 1)

        x.label = "x"
        y.label = "y"
        buf.label = "buf"

        buf2 = m.Float(buf)
        ek.scatter_add(buf2, x, idx1)
        ek.scatter_add(buf2, y, idx2)

        s = ek.dot_async(buf2, buf2)

        if i % 2 == 0:
            ek.enqueue(buf)
        if i // 2 == 0:
            ek.enqueue(x, y)

        ek.traverse(m.Float, reverse=False)

        # Verified against Mathematica
        assert ek.allclose(ek.detach(s), 15.5972)
        assert ek.allclose(ek.grad(s), (25.1667 if i // 2 == 0 else 0) +
                           (17 if i % 2 == 0 else 0))
Exemplo n.º 20
0
def test02_radical_inverse_vectorized(variant_scalar_rgb):
    from mitsuba.core import RadicalInverse

    try:
        from mitsuba.packet_rgb.core.qmc import RadicalInverseP
    except ImportError:
        pytest.skip("packet_rgb mode not enabled")

    v = RadicalInverse()
    v_p = RadicalInverseP()
    for index in range(1024):
        result = v_p.eval_scrambled(index, ek.arange(10, dtype=ek.uint64))
        for i in range(len(result)):
            assert ek.abs(v.eval_scrambled(index, i) - result[i]) < 1e-7
Exemplo n.º 21
0
def test22_scatter_fwd_permute(m):
    x = m.Float(4.0)
    ek.enable_grad(x)

    values_0 = x * ek.linspace(m.Float, 1, 9, 5)
    values_1 = x * ek.linspace(m.Float, 11, 19, 5)

    buf = ek.zero(m.Float, 10)

    idx_0 = ek.arange(m.UInt32, 5)
    idx_1 = ek.arange(m.UInt32, 5) + 5

    ek.scatter(buf, values_0, idx_0, permute=False)
    ek.scatter(buf, values_1, idx_1, permute=False)

    ref = [4.0, 12.0, 20.0, 28.0, 36.0, 44.0, 52.0, 60.0, 68.0, 76.0]
    assert ek.allclose(buf, ref)

    ek.forward(x)
    grad = ek.grad(buf)

    ref_grad = [1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0, 19.0]
    assert ek.allclose(grad, ref_grad)
Exemplo n.º 22
0
def generate_masks(owners, scene):
    #Generate rays that are used for both the binary mask and the distance mask
    sensor = scene.sensors()[0]
    film = sensor.film()
    sampler = sensor.sampler()
    film_size = film.crop_size()
    spp = 3

    # Seed the sampler
    total_sample_count = ek.hprod(film_size) * spp

    if sampler.wavefront_size() != total_sample_count:
        sampler.seed(UInt64.arange(total_sample_count))

    # Enumerate discrete sample & pixel indices, and uniformly sample
    # positions within each pixel.
    pos = ek.arange(UInt32, total_sample_count)
    pos //= spp
    scale = Vector2f(1.0 / film_size[0], 1.0 / film_size[1])
    pos = Vector2f(Float(pos % int(film_size[0])),
                   Float(pos // int(film_size[0])))

    pos += sampler.next_2d()

    # Sample rays starting from the camera sensor
    rays, weights = sensor.sample_ray_differential(time=0,
                                                   sample1=sampler.next_1d(),
                                                   sample2=pos * scale,
                                                   sample3=0)

    #Masks dictionary containas a mask for each params
    masks = dict()
    #For each parameter owner, we generate a mask
    for key in owners:
        scene_dict = dict()
        scene_dict["type"] = "scene"
        scene_dict["camera"] = scene.sensors()[0]
        for item in owners[key]:
            scene_dict[item.id()] = item
        #dummy_scene contains only camera and objects for associated with this owner
        dummy_scene = load_dict(scene_dict)
        surface_interaction = dummy_scene.ray_intersect(rays)

        result = surface_interaction.t + 1.0
        result[~(surface_interaction.t is Infinity)] = 1
        result[~surface_interaction.is_valid()] = 0

        masks[key] = result
    return masks
Exemplo n.º 23
0
def test47_nan_propagation(m):
    for i in range(2):
        x = ek.arange(m.Float, 10)
        ek.enable_grad(x)
        f0 = m.Float(0)
        y = ek.select(x < (20 if i == 0 else 0), x, x * (f0 / f0))
        ek.backward(y)
        g = ek.grad(x)
        if i == 0:
            assert g == 1
        else:
            assert ek.all(ek.isnan(g))

    for i in range(2):
        x = ek.arange(m.Float, 10)
        ek.enable_grad(x)
        f0 = m.Float(0)
        y = ek.select(x < (20 if i == 0 else 0), x, x * (f0 / f0))
        ek.forward(x)
        g = ek.grad(y)
        if i == 0:
            assert g == 1
        else:
            assert ek.all(ek.isnan(g))
Exemplo n.º 24
0
    def load(t, i, offset):
        size = shape[-1 - i]
        stride = strides[-1 - i]

        if i == ndim - 1:
            if type(offset) is int and stride == 1:
                return data
            else:
                i = enoki.arange(enoki.int32_array_t(t), size)
                return t.gather_(data, offset + stride * i, True, False)
        else:
            result = t()
            for j in range(size):
                result[j] = load(t.Value, i + 1, offset + stride * j)
            return result
Exemplo n.º 25
0
def test08_divmod(cname):
    t = get_class(cname)

    index = ek.arange(t, 10000000)
    index[index < len(index) // 2] = -index
    index *= 256203161

    for i in range(1, 100):
        assert index // i == index // ek.opaque(t, i, 1)
        assert index % i == index % ek.opaque(t, i, 1)

    if t.IsSigned:
        for i in range(1, 100):
            assert index // -i == index // ek.opaque(t, -i, 1)
            assert index % -i == index % ek.opaque(t, -i, 1)
Exemplo n.º 26
0
    def sample_ray(
            self,
            time,
            sample1,  # wavelength
            sample2,  # pos
            sample3,  # dir
            active):
        wavelengths, spec_weight = self.m_intensity.sample(
            SurfaceInteraction3f(), ek.arange(sample1), active)
        trafo = self.m_world_transform.eval(ref.time)
        ray = Ray3f(trafo * Point3f(0), warp.square_to_uniform_sphere(sample3),
                    time, wavelengths)

        # print(spec_weight.class_().name())
        return (ray, spec_weight * 4.0 * Pi)
Exemplo n.º 27
0
def test18_slice():
    Float = get_class('enoki.llvm.Float')
    Array2f = get_class('enoki.llvm.Array2f')

    a = ek.arange(Float, 10)
    for i in range(10):
        assert ek.slice(a, i) == i

    a2 = Array2f(a, a * 10)
    for i in range(10):
        assert ek.slice(a2, i)[0] == i
        assert ek.slice(a2, i)[1] == i * 10

    a3 = (a, a * 10)
    for i in range(10):
        assert ek.slice(a3, i)[0] == i
        assert ek.slice(a3, i)[1] == i * 10
Exemplo n.º 28
0
    def sample_ray(
            self,
            time,
            sample1,  # wavelength
            sample2,  # pos
            sample3,  # dir
            active):

        ps = self.m_shape.sample_position(time, sample2, active)
        local = warp.square_to_cosine_hemisphere(sample3)

        si = SurfaceInteraction3f(ps, 0)
        wavelengths, spec_weight = self.m_radiance.sample(
            si, ek.arange(sample1), active)

        ray = Ray3f(ps.p, Frame3f(ps.n).to_world(local), time, wavelengths)
        return (ray, spec_weight * self.m_area_times_pi)
Exemplo n.º 29
0
def test32_atan2(m):
    x = ek.linspace(m.Float, -.8, .8, 10)
    y = m.Float(ek.arange(m.Int, 10) & 1) * 1 - .5
    ek.enable_grad(x, y)
    z = ek.atan2(y, x)
    ek.backward(z)
    assert ek.allclose(
        z,
        m.Float(-2.58299, 2.46468, -2.29744, 2.06075, -1.74674, 1.39486,
                -1.08084, 0.844154, -0.676915, 0.558599))
    assert ek.allclose(
        ek.grad(x),
        m.Float(0.561798, -0.784732, 1.11724, -1.55709, 1.93873, -1.93873,
                1.55709, -1.11724, 0.784732, -0.561798))
    assert ek.allclose(
        ek.grad(y),
        m.Float(-0.898876, -0.976555, -0.993103, -0.83045, -0.344663, 0.344663,
                0.83045, 0.993103, 0.976555, 0.898876))
Exemplo n.º 30
0
def test06_test_collatz(pkg, variant):
    p = get_class(pkg)

    def collatz(value: p.Int):
        counter = p.Int(0)
        loop = p.Loop(value, counter)
        while (loop.cond(ek.neq(value, 1))):
            is_even = ek.eq(value & 1, 0)
            value.assign(ek.select(is_even, value // 2, 3 * value + 1))
            counter += 1
        return value, counter

    value, ctr = collatz(ek.arange(p.Int, 1, 11))
    if variant == 0:
        ek.eval(value, ctr)
    elif variant == 1:
        ek.eval(value)
        ek.eval(ctr)
    elif variant == 2:
        ek.eval(ctr)
        ek.eval(value)

    assert value == p.Int([1] * 10)
    assert ctr == p.Int([0, 1, 7, 2, 5, 8, 16, 3, 19, 6])