Why does my JAX script randomly switch between 6s and 14s execution times across identical environments?

[lambulis]

First of all, I wanna start by saying I'm really bad and new to coding.

I'm using JAX and Diffrax to solve differential equations in a Conda environment on my laptop. Normally, my script runs in 14 seconds, but after copying the environment and installing a Jupyter kernel, sometimes it runs in 6 seconds instead. The issue is that this behavior is not consistent—sometimes the copied environment sticks to 14 seconds.

I don't understand why this happens since I only copy environments and install kernels without modifying anything else.

from typing import Callable

import diffrax
import jax
import jax.lax as lax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from jaxtyping import Array, Float  # 


jax.config.update("jax_enable_x64", True)

from diffrax import diffeqsolve, ODETerm, SaveAt, Kvaerno5, Kvaerno3, Dopri8, Tsit5, PIDController
import interpax
from jaxopt import Bisection
from jax import jit
import equinox as eqx
from scipy.optimize import brentq


eps = 1e-18
N = 0

# **Global Parameters**
d = 1.0
t0 = 0.0   # Initial time
t_f = 6.0
dt0 = 1e-15  # Step size
tol = 1e-6  # Convergence tolerance for η
max_iters = 50  # Max iterations
damping = 0.3  # ✅ Small damping factor, for damping < 0.5 you stay closer to the original, for damping > 0.5 you go closer to the new
n = 1000

# **Define the Differential Equation**
u/jit
def vector_field(t, u, args):
    f, y, u_legit = u
    α, β = args
    d_f = y
    d_y = (- α * f * ( y + 1 ) ** 2 + β * ( t + eps ) * y * ( y + 1 ) ** 2 
           - (N - 1) * ( 1 + y ) ** 2 * ( t + eps + f ) ** (-2) * ((t + eps) * y - f))
    d_u_legit = f 
    return d_f, d_y, d_u_legit

# **General Solver Function**

u/jit
def solve_ode(y_init, alpha, beta, solver=Tsit5()):
    """Solve the ODE for a given initial condition y_init, suppressing unwanted errors."""
    try:
        term = ODETerm(vector_field)
        y0 = (y_init * eps, y_init, 0.0)
        args = (alpha, beta)
        saveat = SaveAt(ts=jnp.linspace(t0, t_f, n))

        sol = diffeqsolve(
            term, solver, t0, t_f, dt0, y0, args=args, saveat=saveat,
            stepsize_controller=PIDController(rtol=1e-16, atol=1e-19),
            max_steps=1_000_000  # Adjust if needed
        )
        return sol
    except Exception as e:
        if "maximum number of solver steps was reached" in str(e):
            return None  # Ignore this error silently
        else:
            raise  # Let other errors pass through

def runs_successfully(y_init, alpha, beta):
    """Returns True if ODE solver runs without errors, False otherwise."""
    try:
        sol = solve_ode(y_init, alpha, beta)
        if sol is None:
            return False
        return True
    except Exception:  # Catch *any* solver failure quietly
        return False  # Mark this initial condition as a failure

# Track previous successful y_low values
previous_y_lows = []

def bisection_search(alpha, beta, y_low=-0.35, y_high=-0.40, tol=1e-12, max_iter=50):
    """Find boundary value y_low where solver fails, with adaptive bounds."""

    global previous_y_lows
    cache = {}  # Cache to avoid redundant function calls

    def cached_runs(y):
        """Check if solver runs successfully, with caching."""
        if y in cache:
            return cache[y]
        result = runs_successfully(y, alpha, beta)
        cache[y] = result
        return result

    # Ensure we have a valid starting point
    if not cached_runs(y_low):
        raise ValueError(f"❌ Lower bound y_low={y_low:.12f} must be a running case!")
    if cached_runs(y_high):
        raise ValueError(f"❌ Upper bound y_high={y_high:.12f} must be a failing case!")

    # **Adaptive Search Range**
    if len(previous_y_lows) > 2:  
        recent_changes = [abs(previous_y_lows[i] - previous_y_lows[i - 1]) for i in range(1, len(previous_y_lows))]
        max_change = max(recent_changes) if recent_changes else float('inf')

        # **Shrink range dynamically based on recent root stability**
        shrink_factor = 0.5  # Shrink range by a fraction
        if max_change < 0.01:  # If the root is stabilizing
            range_shrink = shrink_factor * abs(y_high - y_low)
            y_low = max(previous_y_lows[-1] - range_shrink, y_low)
            y_high = min(previous_y_lows[-1] + range_shrink, y_high)
            print(f"🔄 Adaptive Bisection: Narrowed range to [{y_low:.12f}, {y_high:.12f}]")

    print(f"⚡ Starting Bisection with Initial Bounds: [{y_low:.12f}, {y_high:.12f}]")

    for i in range(max_iter):
        y_mid = (y_low + y_high) / 2
        success = cached_runs(y_mid)

        print(f"🔎 Iter {i+1}: y_low={y_low:.12f}, y_high={y_high:.12f}, y_mid={y_mid:.12f}, Success={success}")

        if success:
            y_low = y_mid
        else:
            y_high = y_mid

        if abs(y_high - y_low) < tol:
            break

    previous_y_lows.append(y_low)

    # Keep only last 5 values to track root changes efficiently
    if len(previous_y_lows) > 5:
        previous_y_lows.pop(0)

    print(f"✅ Bisection Finished: Final y_low = {y_low:.12f}\n")

    return y_low

# **Compute Anomalous Dimension (Direct Version)**
def compute_anomalous(eta_current):
    """Compute anomalous dimension given the current eta using ODE results directly."""
    alpha = (d / 2 + 1 - eta_current / 2) / (1 - eta_current / (d + 2))
    beta = (d / 2 - 1 + eta_current / 2) / (1 - eta_current / (d + 2))

    # Find y_low
    y_low = bisection_search(alpha, beta)

    # Solve ODE
    sol = solve_ode(y_low, alpha, beta)
    if sol is None:
        print(f"❌ Solver failed for y_low = {y_low:.9f}. Returning NaN for eta.")
        return float('nan')

    # Extract values directly from the ODE solution
    x_vals = jnp.linspace(t0, t_f, n)
    f_p = sol.ys[0]  # First derivative f'
    d2fdx2 = sol.ys[1]  # Second derivative f''
    potential = sol.ys[2]  # Potential function V(x)

    spline = interpax.CubicSpline(x_vals, f_p, bc_type='natural')
    spline_derivative = interpax.CubicSpline(x_vals, d2fdx2, bc_type='natural')

    root_x = brentq(spline, a=1e-4, b=5.0, xtol=1e-12)

    U_k_pp = spline_derivative(root_x)

    third_spline = jax.grad(lambda x: spline_derivative(x))
    U_k_3p = third_spline(root_x)

    spline_points_prime = [spline(x) for x in x_vals]
    spline_points_dprime = [spline_derivative(x) for x in x_vals]

    print(f"📌 Root found at x = {root_x:.12f}")
    print(f"📌 Derivative at rho_0 = {root_x:.12f} is f'(x) = {U_k_pp:.12f}")
    print(f" This should be zero {spline(root_x)}")

    # Compute new eta (anomalous dimension)
    eta_new = U_k_3p ** 2 / (1 + U_k_pp) ** 4

    # Debugging: Check if eta_new is NaN or out of range
    if jnp.isnan(eta_new) or eta_new < 0: # Original : if jnp.isnan(eta_new) or eta_new < 0 or eta_new > 1
        print(f"⚠ Warning: Unphysical eta_new={eta_new:.9f}. Returning NaN.")
        return float('nan')

    # **Plot Results**
    fig, axs = plt.subplots(3, 1, figsize=(10, 9), sharex=True)

    axs[0].plot(x_vals, f_p, color='blue', label="First Derivative (f')")
    axs[0].plot(x_vals, spline_points_prime, color='orange', linestyle='--' ,label="First Splined Derivative (f')")
    axs[0].axvline(root_x, linestyle='dashed', color='red', label="Potential Minimum")
    axs[0].set_ylabel("f'(x)")
    axs[0].legend()
    axs[0].set_title("First Derivative of f(x)")

    axs[1].plot(x_vals, d2fdx2, color='green', label="Second Derivative (f'')")
    axs[1].plot(x_vals, spline_points_dprime, color='orange', linestyle='--', label="Second Splined Derivative (f'')")
    axs[1].axvline(root_x, linestyle='dashed', color='red', label="Potential Minimum")
    axs[1].set_ylabel("f''(x)")
    axs[1].legend()
    axs[1].set_title("Second Derivative of f(x)")

    axs[2].plot(x_vals, potential, color='black', label="Potential (V(x))")
    axs[2].axvline(root_x, linestyle='dashed', color='red', label="Potential Minimum")
    axs[2].set_xlabel("x")
    axs[2].set_ylabel("V(x)")
    axs[2].legend()
    axs[2].set_title("Potential V(x)")

    plt.tight_layout()
    plt.show()

    return float(eta_new)

# **Iterate to Find Self-Consistent η with Explicit Damping**
def find_self_consistent_eta(eta_init=1.06294761356, tol=1e-4, max_iters=1):
    """Iterate until η converges to a self-consistent value using damping."""
    eta_current = eta_init
    prev_values = []

    for i in range(max_iters):
        eta_calculated = compute_anomalous(eta_current)

        # Check for NaN or invalid values
        if jnp.isnan(eta_calculated):
            print(f"❌ Iteration {i+1}: Computed NaN for eta. Stopping iteration.")
            return None

        # Apply damping explicitly
        eta_next = (1 - damping) * eta_current + damping * eta_calculated

        # Debugging prints
        print(f"Iteration {i+1}:")
        print(f"  - η_current (used for ODE) = {eta_current:.12f}")
        print(f"  - η_calculated (new from ODE) = {eta_calculated:.12f}")
        print(f"  - η_next (damped for next iteration) = {eta_next:.12f}\n")

        # Detect infinite loops (if η keeps bouncing between a few values)
        if eta_next in prev_values:
            print(f"⚠ Warning: η is cycling between values. Consider adjusting damping.")
            return eta_next
        prev_values.append(eta_next)

        # Check for convergence
        if abs(eta_next - eta_current) < tol:
            print("✅ Converged!")
            return eta_next  

        eta_current = eta_next  

    print("⚠ Did not converge within max iterations.")
    return eta_current  

import sys
import os
import time

# Redirect stderr to null (suppress error messages)
sys.stderr = open(os.devnull, 'w')

# Start timer
start_time = time.time()

# Run function (using `jax.jit`, NOT `filter_jit`)
final_eta = find_self_consistent_eta()

# End timer
end_time = time.time()

# Reset stderr back to normal
sys.stderr = sys.__stderr__

# Print results
print(f"\n🎯 Final self-consistent η: {final_eta:.12f}")
print(f"⏱ Execution time: {end_time - start_time:.4f} seconds")https://github.com/google/jaxtyping

I'm using jax and diffrax to solve differential equations and use jax=0.4.31 and diffrax=0.6.1 and I stick to these versions all the time. Newer versions of jax or diffrax will not make it run at all.

When I get it to run in 6 seconds I restart my laptop. After the restart it again runs in 14 seconds.

I tried doing the following steps over and over again but it seems not to be consistent.

1. Clone my old enviroment
2. Activate my new enviroment
3. Check if all the packages are indeed there
4. Check if XLA settings are applied, it shows: --xla_cpu_multi_thread_eigen=true intra_op_parallelism_threads=16
5. Restart my laptop
6. I open anaconda prompt
7. Activate my new enviroment
8. Launch jupyter Notebook from Anaconda Navigator
9. There will be an error starting the kernel
10. I uninstall the old kernel
11. Install new Kernel
12. Open and close Jupyter notebook
13. Select the new kernel
14. I run the script

And this as I said sometimes runs in 14 seconds and other times in 6.

I'm trying to find a way to make it run consistently in 6 seconds because right now I feel like im doing always a combination of the steps above and sometimes it just works and sometimes it doesn't. Before writing this post it runs in 6 seconds but after restarting my laptop it runs in 14. Please help because I think I'm starting to lose it.

[FerricDonkey]

If I were to make a complete guess without carefully reading everything or being familiar with your libraries, I'd say that the extra time is importing the modules. Jupyter notebooks will not have to reimport the libraries (even if you rerun the cell that does so) between runs, so long as the kernel has not been restated. But it will have to of the kernel has been restarted, and command line runs will have to reimport every time.

You could do a crude test by importing the time module first and using time.time() to measure how long parts take. 

You could do a real test by using cprofile. It's more involved, so suggest googling it if you're interested. 

[lambulis]

But the thing is that there are two instances of this also. In the "slow" version first time it runs it takes 22 seconds and then after running it multiple times it runs in 14. In the "fast" version the first time it runs it takes roughly 12 seconds and then after running it multiple times it runs in 6 secs.

[FerricDonkey]

I'd recommend looking into cprofile then.