const Edge* e;
    if (!send_from_host_node->input_edge(i, &e).ok()) {
      return std::nullopt;
    }

    std::vector<PartialTensorShape> shapes;
    if (!GetNodeAttr(e->src()->attrs(), kXlaInferredShapesAttrName, &shapes)
             .ok()) {
      return std::nullopt;
    }

    const PartialTensorShape shape = shapes[e->src_output()];
    if (!shape.IsFullyDefined()) {
      return std::nullopt;
    }