Change KerasTensor.set_shape back to mutate the nested TensorSpec in-place.
@@ -636,7 +636,17 @@ def keras_tensor_from_type_spec(type_spec, name=None):
def type_spec_with_shape(spec, shape):
"""Returns a copy of TypeSpec `spec` with its shape set to `shape`."""
if isinstance(spec, tf.TensorSpec):
- return tf.TensorSpec(shape, spec.dtype, spec.name)
+ # pylint: disable=protected-access
+ # TODO(b/203201161) Figure out why mutation is needed here, and remove it.
+ # (TensorSpec objects should be immutable; and we should not be modifying
+ # private fields.)
+ shape = tf.TensorShape(shape)
+ spec._shape = shape
+ if shape.rank is None:
+ spec._shape_tuple = None
+ else:
+ spec._shape_tuple = tuple(shape.as_list())
+ return spec
elif isinstance(spec, tf.RaggedTensorSpec):
return tf.RaggedTensorSpec(shape, spec.dtype, spec.ragged_rank,
spec.row_splits_dtype,
