bool InferShapeForXlaHostComputeMlir(_XlaHostComputeMlirOp op);

  // Infers the shape of function attached to XlaHostCompute.
  // Returns true if a return type was changed.
  bool InferShapeForFunctionAttachedToXlaHostCompute(XlaHostComputeOp op);

  // Infers the shape for MapDatasetOp and its associated function. Returns
  // whether either op or function type was changed.