code: def calculate_allocated_bytes(
      func_or_mod: Union[PrimFunc, IRModule]
  ) -> Union[Dict[str, int], Dict[str, Dict[str, int]]]:
      """Calculate allocated memory per memory scope required by TIR PrimFuncs.
  
      Parameters
      ----------
      func_or_mod: Union[PrimFunc, IRModule]
          The function or module to be detected. If a module is passed, allocated
          memory is calculated for all PrimFuncs inside the module
  
      Returns
      -------
      result : Union[Dict[str, int], Dict[str, Dict[str, int]]]
          Allocated memory size per scope in bytes for each function in the IRModule returned as a
          dict with function names as keys and a dict of allocated sizes as values. If a single
          PrimFunc is passed, the function name is returned as "main"
      """
      if not isinstance(func_or_mod, (PrimFunc, IRModule)):
          raise TypeError(
              f"Expected argument to be PrimFunc or IRModule, but received {type(func_or_mod)}"
          )
      return _ffi_api.calculate_allocated_bytes(func_or_mod)
  
imports: from typing import Dict, List, Optional, Union
  import tvm
  from tvm import Object
  from tvm.ir import IRModule
  from tvm.tir.expr import Var
  from tvm.tir.stmt import Block, BufferRegion, PrimExpr
  from ..None import Buffer, Stmt
  from ..function import PrimFunc
  from .None import _ffi_api
label_desc: Use image node_id 9 for calling a global function with example usage: calculate_allocated_bytes(func_or_mod) and returns: _ffi_api