code: def recompute_grads_w_kwargs(
      fn, use_entire_scope=False, train_vars=None, name=None
  ):
      def inner_recompute_grads_w_kwargs(*args, **kwargs):
          # Force an activation cache just before it enters the function, meaning output of 1 never == input of the other.
          args = [
              tf.identity(a) if isinstance(a, tf.Tensor) else a
              for a in args
          ]
          kwargs = {
              k: tf.identity(v) if isinstance(v, tf.Tensor) else v
              for k, v in kwargs.items()
          }
          output_keywords = None
          output_local = None
  
          def remapped_fn(*_):
              nonlocal output_keywords
              nonlocal output_local
              out = fn(*args, **kwargs)
              if isinstance(out, dict):
                  output_keywords = []
                  output_local = dict()
                  for k, v in out.items():
                      if isinstance(v, tf.Tensor):
                          output_keywords.append(k)
                      else:
                          output_local[k] = v
                  return [out[k] for k in output_keywords]
              return out
  
          tensor_args_kwargs = [
              a
              for a in list(args) + list(kwargs.values())
              if isinstance(a, tf.Tensor)
          ]
          out = _recompute_grad(
              remapped_fn,
              tensor_args_kwargs,
              use_entire_scope=use_entire_scope,
              train_vars=train_vars,
          )
          if output_keywords is not None:
              return {
                  **{k: v for k, v in zip(output_keywords, out)},
                  **output_local,
              }
          return out
  
      return inner_recompute_grads_w_kwargs
  
file_path: finetune/finetune/optimizers/recompute_grads.py
imports: import functools
  import tensorflow
  from tensorflow.python.framework import function, ops
  from finetune.util.optimize_loss import is_numerical_tensor
label_desc: Use image node_id 6 for calling a global function with example usage: recompute_grads_w_kwargs(fn, use_entire_scope, train_vars, name) and returns: inner_recompute_grads_w_kwargs, out, out, dict, unknown