def _tf_f(*args, **kwargs):
    """Accelerated function with tensor inputs/outputs."""
    np_args = _tf_to_np(args)
    kwargs = {k: _tf_to_np(v) for k, v in kwargs.items()}
    if xla_forced_compile:
      # Use list for mutability
      output_is_list = [False]
      output_is_empty = [False]
      output_structure = [None]
      def recorder(args, kwargs, res):
        del args, kwargs
        # Workaround b/121383831
        output_is_list[0] = isinstance(res, list)
        # If outputs are empty, xla.compile returns an `Operation`, which we
        # don't want.
        if tf.nest.flatten(res):
          output_is_empty[0] = False
          output_structure[0] = None
        else:
          output_is_empty[0] = True
          # Without deepcopy, xla.compile will change output_structure[0] to a
          # list of `Operation`.
          output_structure[0] = copy.deepcopy(res)
        return res
      f_ = _record_result_type(recorder, f)
      np_out = tf.xla.experimental.compile(lambda: f_(*np_args, **kwargs))
      # Workaround b/121383831
      if output_is_empty[0]:
        np_out = output_structure[0]
      elif (isinstance(np_out, list) and len(np_out) == 1 and
            not output_is_list[0]):
        np_out = np_out[0]
    else:
      np_out = f(*np_args, **kwargs)
    return np_out
Accelerated function with tensor inputs/outputs.
trax/tf_numpy/extensions/extensions.py