code: def cond_batch_rule(interpreter, pred, true_fn, false_fn, inputs):
      assert isinstance(
          inputs, (list, tuple)
      ), "Cond inputs must be a list or tuple of tensors"
      assert all(
          isinstance(i, torch.Tensor) for i in inputs
      ), "Cond inputs must be a list of tensors"
  
      pred_ = get_unwrapped(pred) if is_batchedtensor(pred) else pred
  
      # unbatched tensors are not vmapped
      tensors, in_dims = zip(
          *[
              (get_unwrapped(t), maybe_get_bdim(t))
              if is_batchedtensor(t)
              else (t, None)
              for t in inputs
          ]
      )
  
      if is_batchedtensor(pred):
          # prepend "pred" and vmap everything
          tensors = (pred_,) + tensors
          in_dims = (0,) + in_dims
  
          def fn(p, *args):
              t = true_fn(*args)
              f = false_fn(*args)
              return torch.where(p, t[0], f[0])
  
          with interpreter.lower():
              result = torch.vmap(fn, in_dims=in_dims)(*tensors)
  
      else:
          # predicate is known at this stage and it is a boolean expression or a
          # tensor with one element.
          true_fn = torch.vmap(true_fn, in_dims=in_dims)
          false_fn = torch.vmap(false_fn, in_dims=in_dims)
  
          with interpreter.lower():
              result = cond_op(pred, true_fn, false_fn, tensors)
  
      if not isinstance(result, tuple):
          result = (result,)
      lvl = interpreter.level()
      return tuple([_add_batch_dim(r, 0, lvl) for r in result])
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  import torch
  import torch._subclasses.functional_tensor
  import torch.fx.traceback
  import torch.utils._pytree
  from torch._C import DispatchKey
  from torch._C._functorch import _add_batch_dim, get_unwrapped, is_batchedtensor, maybe_get_bdim
  from torch._functorch.utils import exposed_in
  from torch._higher_order_ops.utils import autograd_not_implemented
  from torch._ops import HigherOrderOperator
  from torch._subclasses.fake_tensor import FakeTensorMode
  from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing, make_fx, ProxyTorchDispatchMode, track_tensor_tree
  from torch.fx.passes.shape_prop import _extract_tensor_metadata
  from torch.multiprocessing.reductions import StorageWeakRef
  from torch.utils._python_dispatch import _get_current_dispatch_mode
label_desc: Use image node_id 11 for calling a global function with example usage: cond_batch_rule(interpreter, pred, true_fn, false_fn, inputs) and returns: tuple, torch