def tf_f(*tf_args):
    """A wrapper for `f` that takes/returns tensors."""
    np_args = _tf_to_np(tf_args)
    np_out = f(*np_args)
    return np_out
A wrapper for `f` that takes/returns tensors.
trax/tf_numpy/extensions/extensions.py