def METHOD_NAME(tf_fun: Callable, arg, *, result_shape):
  """The simplest implementation of calling to TF, without AD support.
  We must use hcb.call because the TF invocation must happen outside the
  JAX staged computation."""
  def tf_to_numpy(t):
    # Turn the Tensor to NumPy array without copying.
    return np.asarray(memoryview(t)) if isinstance(t, tf.Tensor) else t
  return hcb.call(lambda arg: tf.nest.map_structure(tf_to_numpy,
                                                    tf_fun(arg)),
                  arg, result_shape=result_shape)