def testIssue711(self):

    def harmonic_bond(conf, params):
        return jnp.sum((conf * params))

    def minimize_structure(test_params):
        energy_fn = partial(harmonic_bond, params=test_params)

        def apply_carry(carry, _):
            (i, x) = carry
            new_x = (x - (0.1 * api.grad(energy_fn)(x)))
            new_carry = ((i + 1), new_x)
            return (new_carry, _)
        x0 = jnp.array([1.0, 2.0, 3.0])
        (carry_final, _) = lax.scan(apply_carry, (0, x0), jnp.zeros((75, 0)))
        (_, x_final) = carry_final
        return x_final
    initial_params = 0.5
    minimize_structure(initial_params)

    def loss(test_params):
        x_final = minimize_structure(test_params)
        return jnp.sum(jnp.sin((1.0 - x_final)))
    api.grad(loss)(0.25)
def testIssue711(self):

    def harmonic_bond(conf, params):
        return jnp.sum((conf * params))

    def minimize_structure(test_params):
        energy_fn = partial(harmonic_bond, params=test_params)

        def apply_carry(carry, _):
            (i, x) = carry
            new_x = (x - (0.1 * api.grad(energy_fn)(x)))
            new_carry = ((i + 1), new_x)
            return (new_carry, _)
        x0 = jnp.array([1.0, 2.0, 3.0])
        (carry_final, _) = lax.scan(apply_carry, (0, x0), jnp.zeros((75, 0)))
        (_, x_final) = carry_final
        return x_final
    initial_params = 0.5
    minimize_structure(initial_params)

    def loss(test_params):
        x_final = minimize_structure(test_params)
        return jnp.sum(jnp.sin((1.0 - x_final)))
    api.grad(loss)(0.25)