code: def all_bdims(*shapes):
      bdims = (
          itertools.chain(
              [cast(Union[int, None], None)], range(len(shape) + 1)
          )
          for shape in shapes
      )
      return (
          t
          for t in itertools.product(*bdims)
          if not all(e is None for e in t)
      )
  
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 4 for calling a global function with example usage: all_bdims() and returns: unknown