Type potential_refined_type);

  // Infers the shape from a (Stateful)PartitionedCall operation by looking up
  // the called function and propagating the return type.
  bool InferShapeForCall(CallOpInterface call_op);

  bool InferShapeForCast(Operation* op);

  bool InferShapeForRestore(Operation* op);

  // Infers the shape IfOp outputs based on the shapes of the then and else