code: def mixed_checkpoint(func, inputs: dict, params, flag):
      """
      Evaluate a function without caching intermediate activations, allowing for
      reduced memory at the expense of extra compute in the backward pass. This differs from the original checkpoint function
      borrowed from https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py in that
      it also works with non-tensor inputs
      :param func: the function to evaluate.
      :param inputs: the argument dictionary to pass to `func`.
      :param params: a sequence of parameters `func` depends on but does not
                     explicitly take as arguments.
      :param flag: if False, disable gradient checkpointing.
      """
      if flag:
          tensor_keys = [
              key
              for key in inputs
              if isinstance(inputs[key], torch.Tensor)
          ]
          tensor_inputs = [
              inputs[key]
              for key in inputs
              if isinstance(inputs[key], torch.Tensor)
          ]
          non_tensor_keys = [
              key
              for key in inputs
              if not isinstance(inputs[key], torch.Tensor)
          ]
          non_tensor_inputs = [
              inputs[key]
              for key in inputs
              if not isinstance(inputs[key], torch.Tensor)
          ]
          args = (
              tuple(tensor_inputs)
              + tuple(non_tensor_inputs)
              + tuple(params)
          )
          return MixedCheckpointFunction.apply(
              func,
              len(tensor_inputs),
              len(non_tensor_inputs),
              tensor_keys,
              non_tensor_keys,
              *args,
          )
      else:
          return func(**inputs)
  
imports: import math
  from typing import Optional
  import torch
  import torch.nn
  from einops import rearrange, repeat
label_desc: Use image node_id 3 for calling a global function with example usage: mixed_checkpoint(func, inputs, params, flag) and returns: MixedCheckpointFunction, func