code: def add_bdim(bdim_size, bdim, shape):
      shape = list(shape)
      if bdim is not None:
          shape.insert(bdim, bdim_size)
      return tuple(shape)
  
file_path: jax/jax/_src/internal_test_util/lax_test_util.py
imports: from __future__ import annotations
  import collections
  import itertools
  from typing import Union, cast
  from jax import lax
  from jax._src import dtypes
  from jax._src import test_util
  from jax._src.util import safe_map, safe_zip
  import numpy
  from jax import config
label_desc: Use image node_id 5 for calling a global function with example usage: add_bdim(bdim_size, bdim, shape) and returns: tuple