code: def testCustomJVPLeak(self):
      # https://github.com/google/jax/issues/8171
      @jax.jit
      def fwd():
          a = jnp.array(1.0)
  
          def f(hx, _):
              hx = jax.nn.sigmoid(hx + a)
              return hx, None
  
          hx = jnp.array(0.0)
          jax.lax.scan(f, hx, None, length=2)
  
      with jax.checking_leaks():
          fwd()
  
imports: import collections
  from functools import partial
  import itertools
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import scipy.stats
  from jax._src import config
  from jax._src import core
  from jax._src import test_util
  from jax._src import ad_checkpoint
  from jax.test_util import check_grads
  from jax import nn
  from jax import random
  import jax
  import jax.numpy
label_desc: Use image node_id 36 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testCustomJVPLeak() and returns: hx,  None