code: def compute_expected_grads(
      op,
      args,
      kwargs,
      output_process_fn_grad=None,
      gradcheck_wrapper=None,
  ):
      if gradcheck_wrapper is None:
          results = op(*args, **kwargs)
      else:
          results = gradcheck_wrapper(op, *args, **kwargs)
  
      if output_process_fn_grad is not None:
          results = output_process_fn_grad(results)
  
      flat_results = pytree.tree_leaves(results)
      flat_results = [
          r for r in flat_results if isinstance(r, torch.Tensor)
      ]
      flat_diff_results = [r for r in flat_results if r.requires_grad]
      assert len(flat_diff_results) > 0
  
      grads = [
          torch.ones(r.shape, device=r.device, dtype=r.dtype)
          for r in flat_diff_results
      ]
      leaf_tensors = gather_leaf_tensors(args, kwargs)
      assert len(leaf_tensors) > 0
      return torch.autograd.grad(
          flat_diff_results,
          leaf_tensors,
          grads,
          allow_unused=True,
          retain_graph=True,
      )
  
imports: import torch
  from torch import Tensor
  import itertools
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
  from torch.utils import _pytree
  from functools import partial
  from torch.utils._mode_utils import no_dispatch, all_same_mode
  import torch.autograd.forward_ad
  from typing import Callable
  import re
label_desc: Use image node_id 16 for calling a global function with example usage: compute_expected_grads(op, args, kwargs, output_process_fn_grad, gradcheck_wrapper) and returns: torch