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
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