提交 d05c6ef9 编辑于 作者: Keith Rush's avatar Keith Rush 提交者: tensorflow-copybara
浏览文件

Adjust key computation for TF functions to use determinstic proto serialization.

Also serializes TFF types and uses the same path, since TFF TensorType strings dont currently differentiate between TensorShape(None) and scalar tensors.

PiperOrigin-RevId: 346613971
上级 8d719b06
......@@ -294,8 +294,9 @@ def _to_computation_internal_rep(*, value: pb.Computation,
type_spec: computation_types.StructType,
device: tf.config.LogicalDevice):
"""Converts a `pb.Computation` to a `tf.function`."""
key = (value.SerializeToString(), str(type_spec),
device.name if device else None)
key = (value.SerializeToString(deterministic=True),
type_serialization.serialize_type(type_spec).SerializeToString(
deterministic=True), device.name if device else None)
cached_fn = tf_function_cache.get(key)
if cached_fn is not None:
return cached_fn
......
Supports Markdown
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册