code: def reconstruct_cond_batch(
      c: list[list[ScheduledPromptConditioning]], current_step
  ):
      param = c[0][0].cond
      is_dict = isinstance(param, dict)
  
      if is_dict:
          dict_cond = param
          res = {
              k: torch.zeros(
                  (len(c),) + param.shape,
                  device=param.device,
                  dtype=param.dtype,
              )
              for k, param in dict_cond.items()
          }
          res = DictWithShape(
              res, (len(c),) + dict_cond["crossattn"].shape
          )
      else:
          res = torch.zeros(
              (len(c),) + param.shape,
              device=param.device,
              dtype=param.dtype,
          )
  
      for i, cond_schedule in enumerate(c):
          target_index = 0
          for current, entry in enumerate(cond_schedule):
              if current_step <= entry.end_at_step:
                  target_index = current
                  break
  
          if is_dict:
              for k, param in cond_schedule[target_index].cond.items():
                  res[k][i] = param
          else:
              res[i] = cond_schedule[target_index].cond
  
      return res
  
imports: from __future__ import annotations
  import re
  from collections import namedtuple
  import lark
label_desc: Use image node_id 5 for calling a global function with example usage: reconstruct_cond_batch(c, current_step) and returns: res