from tensorflow.python.framework import ops

  # TODO(mrry): Move this to where it is used, so we can get rid of this op
  #   wrapper?
  if set_shape:
    ret.set_shape(shape)
  return ret


# NOTE(mrry): Shapes are conditionally set in the Python wrapper.
ops.RegisterShape("Variable")(common_shapes.unknown_shape)


@ops.RegisterShape("TemporaryVariable")
def _TemporaryVariableShape(op):
  """Shape function for the TemporaryVariable op."""
  shape = tensor_util.TensorShapeProtoToList(op.get_attr("shape"))
  return [tensor_shape.TensorShape(shape)]
