return success();
  }
};

void SessionInitializerOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<OptimizeSessionInitializerPattern>(context);
}

SmallVector<StringRef, 2> GetSessionInitializerExportedName(ModuleOp op) {
  auto session_initializer_op = GetSessionInitializerOp(op);

        "tf_device.cluster"() ({
          %group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
          %arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
          "tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
          tf_device.return
        }) : () -> ()
        tf_device.return
      }
      return
    }

        TypeID::get<TF::XlaSpmdFullToShardShapeOp>(),
        TypeID::get<TF::XlaSpmdShardToFullShapeOp>(),
        TypeID::get<TF::XlaSvdOp>(),
    };

    // Add the ops from the TPUEmbeddingOpsRegistry.
    for (auto op_type_id :
         TF::TPUEmbeddingOpsRegistry::Global().GetOpsTypeIds()) {
      ops_set->insert(op_type_id);
    }
    return ops_set;
  }();

               std::back_inserter(result), [](const BufferInfo& buffer_info) {
                 return buffer_info.is_temp_buffer();
               });
  return result;
}

// Add (from,to) rewrite pairs based on the given shape.  These rewrite pairs
// are used to generate methods for args and results.
Status AddRewritesForShape(int i, const xla::Shape& shape,

// If GCS_APPEND_MODE=compose then instead the new data is uploaded to a
// temporary object and composed with the original object. This is disabled by
// default as the multiple API calls required add a risk of stranding temporary
// objects.
constexpr char kComposeAppend[] = "compose";

// We can cast `google::cloud::StatusCode` to `TF_Code` because they have the
// same integer values. See