Unexpected memory leak during the forward pass of a gradient calculation

[fishjojo]

Dear developers,

I have a simple test as follows

import os, sys
import jax
from jax import numpy as np
jax.config.update("jax_enable_x64", True)

def current_memory():
    '''Get the current RSS in MB'''
    assert sys.platform.startswith('linux')
    with open("/proc/%s/statm" % os.getpid()) as f:
        vms, rss = [int(x)*os.sysconf("SC_PAGE_SIZE") for x in f.readline().split()[:2]]
    return int(rss/1e6)

@jax.custom_vjp
def get_jk(eri, dm):
    vj = vk = dm
    return vj, vk

def get_jk_fwd(eri, dm):
    vj, vk = get_jk(eri, dm)
    return (vj, vk), None

def get_jk_bwd(res, ybar):
    return None, None

get_jk.defvjp(get_jk_fwd, get_jk_bwd)

n = 500
h1 = dm0 = np.ones((n,n))  # 2MB
eri = np.ones((n*6, n*(n+1)//2)) # 3GB

def energy(h1):
    etot = 0
    dm = dm0
    for i in range(10):
        vj, vk = get_jk(eri, dm)
        vhf = vj - .5 * vk + h1
        dm += h1
        etot += np.trace(np.dot(vhf, dm))
        print(f'RSS: {current_memory()} MB')
    return etot

e, grad = jax.value_and_grad(energy)(h1)

and the output looks like

RSS: 3168 MB
RSS: 6182 MB
RSS: 9194 MB
RSS: 12205 MB
RSS: 15218 MB
RSS: 18232 MB
RSS: 21245 MB
RSS: 24256 MB
RSS: 27267 MB
RSS: 30278 MB

As you can see, for each iteration in the forward pass, the memory usage increases about 3GB, which is the size of the eri tensor. However, what I expected is a roughly constant memory usage of 3GB, since I didn't save any residual quantities during the forward pass. In the real application, the function get_jk calls foreign functions, so I can't use jit compile here. Please help me understand what causes the memory leak.
The test was run on a red hat linux x86-64 server and Jax v0.4.25 was used.