Fortify access to ScopedTFGraph with a context manager.

Such that the backing C object it is never destroyed in the middle of a
python function that uses it. It is unclear which functions that uses _c_graph can be called on the problematic 'use-during-deletion' path, thus I changed them all to use the fortified .get() API mechanically.

While at it, I added a tensor_util API to reduce number of protected access violations caused by TF_TryEvaluateConstant. It is unclear how the C++ evaluate constant relates to the _ConstantValue implemented in tensor_util.py. At least the names are different, and they are now in the same file, hopefully making a future consolidation easier.

PiperOrigin-RevId: 455227569

Author: Googler

tensorflow_probability/python/internal/prefer_static.py

@@ -25,15 +25,9 @@
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.internal.backend import numpy as nptf
 
-# Try catch required to avoid breaking Probability opensource presubmits.
-# TODO(amitpatankar): Remove this once tf-nightly has latest code.
-# pylint: disable=g-import-not-at-top
-try:
-  from tensorflow.python.client import pywrap_tf_session as c_api  # pylint: disable=g-direct-tensorflow-import
-except ImportError:
-  from tensorflow.python import pywrap_tensorflow as c_api  # pylint: disable=g-direct-tensorflow-import
-
+from tensorflow.python.client import pywrap_tf_session as c_api  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.framework import ops  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.framework import tensor_util  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import control_flow_ops  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.util import tf_inspect  # pylint: disable=g-direct-tensorflow-import
 
@@ -114,14 +108,16 @@ def _get_static_value(pred):
   if tf.is_tensor(pred):
     pred_value = tf.get_static_value(tf.convert_to_tensor(pred))
 
-    # TODO(jamieas): remove the dependency on `pywrap_tensorflow`.
     # Explicitly check for ops.Tensor, to avoid an AttributeError
     # when requesting `KerasTensor.graph`.
-    # pylint: disable=protected-access
     if pred_value is None and isinstance(pred, ops.Tensor):
-      pred_value = c_api.TF_TryEvaluateConstant_wrapper(pred.graph._c_graph,
-                                                        pred._as_tf_output())
-    # pylint: enable=protected-access
+      if hasattr(tensor_util, 'try_evaluate_constant'):
+        pred_value = tensor_util.try_evaluate_constant(pred)
+      else:
+        # TODO(feyu): remove this branch after try_evaluate_constant is in
+        # tf-nightly.
+        pred_value = c_api.TF_TryEvaluateConstant_wrapper(
+            pred.graph._c_graph, pred._as_tf_output())  # pylint: disable=protected-access
     return pred_value
   return pred