Function GetShapeValue

serving/processor/framework/graph_optimizer.cc:730–753

for KvVarHandleOp and MutableHashTableOfTensorsV2

// for KvVarHandleOp and MutableHashTableOfTensorsV2
Status GetShapeValue(Node* node, int* dim) {
  AttrValue* dim_len_value =
      const_cast<AttrValue*>(node->attrs().Find("shape"));
  if (!dim_len_value) {
    dim_len_value = const_cast<AttrValue*>(node->attrs().Find("value_shape"));
  }

  Status s = AttrValueHasType(*dim_len_value, "shape");
  if (!s.ok()) {
    return tensorflow::errors::Internal(
        "Miss shape attr in the KvVarHandleOp, ",
        node->DebugString());
  }

  if (dim_len_value->shape().dim().size() == -1) {
    return tensorflow::errors::Internal(
        "Dim value of the shape in the KvVarHandleOp is unknown, ",
        node->DebugString());
  }

  *dim = dim_len_value->shape().dim(0).size();

  return Status::OK();
}

bool IsKvOps(const Node* node) {
  return node->op_def().name() == "KvResourceGather" ||