code: def test_inline_jit_location_uses_callee_location(self):
      # 'f' should be inlined into both 'g' and 'h', using the source line
      # information of the call site. In particular, the source line information
      # of 'h' should not refer to the source information of 'g'.
      @partial(jax.jit, inline=True)
      def f(x):
          return lax.add(x, 3)
  
      def g(x):
          return lax.add(f(x), 4)
  
      def h(x):
          return lax.add(f(x), 5)
  
      for fn in (g, h):
          lines, fn_startline = inspect.getsourcelines(fn)
          fn_endline = fn_startline + len(lines)
          jaxpr = jax.make_jaxpr(fn)(2)
          for eqn in jaxpr.eqns:
              frame = source_info_util.user_frame(eqn.source_info)
              assert frame is not None, eqn
              self.assertLessEqual(fn_startline, frame.start_line)
              self.assertLessEqual(frame.end_line, fn_endline)
  
file_path: jax/tests/source_info_test.py
imports: from functools import partial
  import inspect
  from absl.testing import absltest
  import jax
  from jax import lax
  from jax import config
  from jax._src import source_info_util
  from jax._src import test_util
label_desc: Use image node_id 1 for calling the SourceInfoTest obj's underlying member method code with example usage: obj.test_inline_jit_location_uses_callee_location() and returns: lax, lax, lax