def _record_result_type(recorder, f):
  """A decorator that records some information about the function.

  Args:
    recorder: a function of signature `(args, kwargs, res) -> res`.
    f: the original function.

  Returns:
    A transformed function that calls the original function and then the
    recorder afterwards.
  """
  def wrapper(*args, **kwargs):
    res = f(*args, **kwargs)
    res = recorder(args, kwargs, res)
    return res

  return wrapper
A decorator that records some information about the function.

  Args:
    recorder: a function of signature `(args, kwargs, res) -> res`.
    f: the original function.

  Returns:
    A transformed function that calls the original function and then the
    recorder afterwards.
  
trax/tf_numpy/extensions/extensions.py