Diminishing speed-up with increasing problem size using lax.scan()

[n1ck94]

Hi jax team,
I am trying to solve the time-dependent heat equation as fast as possible. Below is a simplified version of my program.

numGPs = 25
tdur = 100
dt = 0.025
dx = 0.05
kappa = 10.
rhoCp = 2000.
eta = 25.

np.random.seed(10)
T0 = np.ones(numGPs)*25.
numSteps = int(tdur/dt)
TAmb = np.random.random(numSteps)*2 + 25.

def get_Tn(carry, x):
    Tn_, dt, kappa, rhoCp, eta, dx = carry
    TAmbn = x
    qL = -eta*(TAmbn - Tn_[0])
    qX = kappa*(Tn_[1:] - Tn_[:-1])/dx
    qR = eta*(TAmbn - Tn_[-1])
    Tn = Tn_.at[1:-1].add(dt/rhoCp*(qX[1:] - qX[:-1])/dx)
    Tn = Tn.at[0].add(dt/rhoCp*2*(qX[0] - qL)/dx)
    Tn = Tn.at[-1].add(dt/rhoCp*2*(qR - qX[-1])/dx)
    return (Tn, dt, kappa, rhoCp, eta, dx), Tn

def TimeStepping(T0, dt, kappa, rhoCp, eta, dx, TAmb):
    _, T = jax.lax.scan(get_Tn, (T0, dt, kappa, rhoCp, eta, dx), TAmb)
    return T

numTests = 3
test_record = []
for i in range(numTests):
    t0 = time.time()
    T = TimeStepping(T0, dt, kappa, rhoCp, eta, dx, TAmb)
    test_record.append(time.time()-t0)

fT_Value_jit = jax.jit(TimeStepping)
fT_Value_jit(T0, dt, kappa, rhoCp, eta, dx, TAmb).block_until_ready()
test_record_jit = []
for i in range(numTests):
    t0 = time.time()
    T = fT_Value_jit(T0, dt, kappa, rhoCp, eta, dx, TAmb)
    test_record_jit.append(time.time()-t0)
print(test_record)
print(test_record_jit)
print(np.average(test_record) / np.average(test_record_jit))

The following is a summary of the speed-up for the jitted function with increasing problem size (i.e. increasing number of nodes/gauss points numGPs in the domain):

• numGPs=25, 110x
• numGPs=250, 11x
• numGPs=2500, 1.9x
• numGPs=25000, 1.1x

I would like to know if the diminishing speed-up is the expected behaviour or if I am doing something silly. If this is the expected behaviour, would anyone have any suggestions for improving the speed-up? I think it would be great to use vmap(), however, as the temperature at time step $t^n$ is a function of the temperature at the previous time step $t^{n-1}$, I'm not sure how I could implement it.

I have opted for scan() over fori_loop() for example as I ultimately need the gradients. Thus far, I have only been running on CPU. I'm running python 3.11, jax 0.4.1, jaxlib 0.4.1.

Thanks!