code: def _extract_batch_size(
      batch: BType,
  ) -> Generator[Optional[int], None, None]:
      if isinstance(batch, Tensor):
          if batch.ndim == 0:
              yield 1
          else:
              yield batch.size(0)
      elif isinstance(batch, (Iterable, Mapping)) and not isinstance(
          batch, str
      ):
          if isinstance(batch, Mapping):
              batch = batch.values()
  
          for sample in batch:
              yield from _extract_batch_size(sample)
      elif is_dataclass_instance(batch):
          for field in fields(batch):  # type: ignore[arg-type]
              yield from _extract_batch_size(getattr(batch, field.name))
      else:
          yield None
  
imports: import inspect
  from dataclasses import fields
  from typing import Any, Dict, Generator, Iterable, Mapping, Optional, Sized, Tuple, Union
  import torch
  from lightning_utilities.core.apply_func import is_dataclass_instance
  from torch import Tensor
  from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler, Sampler, SequentialSampler
  from typing_extensions import TypeGuard
  import lightning.pytorch
  from lightning.fabric.utilities.data import _reinstantiate_wrapped_cls, _replace_value_in_saved_args, has_iterable_dataset, sized_len
  from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper
  from lightning.pytorch.trainer.states import RunningStage
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.rank_zero import WarningCache, rank_zero_warn
label_desc: Use image node_id 1 for calling a global function with example usage: _extract_batch_size(batch) without return types