class_docstr: Implements the TF backend ops for a PC auto-batching VM.
docstr: Method exists solely to be stubbed, i.e. for defun + XLA compile.
docstr_tok: ["Method","exists","solely","to","be","stubbed",",","i.e",".","for","defun","+","XLA","compile","."]
code: def wrap_straightline_callable(self, f):
      """Method exists solely to be stubbed, i.e. for defun + XLA compile."""
      if self._basic_block_xla_device is None:
          return f
  
      @tf.function
      def _f(*args):
          with tf.device(self._basic_block_xla_device):
              return xla.compile_nested_output(
                  f, tf.xla.experimental.compile
              )(*args)
  
      def _ensure_regvars_initialized(t):
          if isinstance(t, RegisterTensorFlowVariable):
              return t.ensure_initialized()
          return t
  
      def _init_f(env_dict, *args):
          """A RegisterTensorFlowVariable-initializing wrapper of `_f`."""
          # We ensure RegisterTensorFlowVariable instances have a Tensor value when
          # using XLA and/or defun. Otherwise, we will trigger cache misses on the
          # tfe.defun or get issues around "Cannot convert object of type [dtype] to
          # a Tensor" (XLA). This corresponds with the optimization in
          # `create_variable` conditioned on Eager & VariableAllocation.REGISTER.
          env_dict = dict(
              {
                  k: instructions.pattern_map(
                      _ensure_regvars_initialized,
                      v,
                      leaf_type=RegisterTensorFlowVariable,
                  )
                  for k, v in six.iteritems(env_dict)
              }
          )
          return _f(env_dict, *args)
  
      return _init_f
  
imports: import collections
  import contextlib
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import xla
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 23 for calling the TensorFlowBackend obj's underlying member method code with example usage: obj.wrap_straightline_callable(f) and returns: _init_f, f, t, _f, xla, t