xla::ShapeUtil::DeviceShapeToHostShape(device_shape);

    auto resource_var_it = resource_vars.find(arg_num);
    bool is_resource_variable = resource_var_it != resource_vars.end();
    bool is_updated_resource_variable =
        is_resource_variable &&
        absl::c_any_of(compilation_result->resource_updates,
                       [&](const XlaCompiler::ResourceUpdate& update) {