Method session

tensorflow/compiler/tests/xla_test.py:204–224

Custom implementation of session() for XLA tests. We override the standard Tensorflow session() since it is too specific to CPU and GPU tests. In particular, we want to disable soft placement and explicitly assign ops to devices under test. Yields: A session to use when runni

(self)

  @contextlib.contextmanager
  def session(self):
    """Custom implementation of session() for XLA tests.

    We override the standard Tensorflow session() since it is too
    specific to CPU and GPU tests. In particular, we want to disable soft
    placement and explicitly assign ops to devices under test.

    Yields:
      A session to use when running a test case.
    """
    graph = ops.Graph()
    config = context.context().config

    # Grappler can constant fold TensorListFromTensor ops into DT_VARIANT
    # constants which XLA does not understand.  So disable constant folding in
    # these tests.
    config.graph_options.rewrite_options.constant_folding = (
        rewriter_config_pb2.RewriterConfig.OFF)
    with session.Session(
        graph=graph, config=config) as sess, graph.as_default():
      yield sess

  def cached_session(self):
    raise NotImplementedError(