message: Add attribute "num_distributed_variables" to _TPUReplicate op, in order to distinguish per-replica inputs and distributed variable inputs.
diff: @@ -108,6 +108,7 @@ REGISTER_OP("_TPUReplicate")
  .Attr("padding_map: list(string) = []")
  .Attr("step_marker_location: string = \"STEP_MARK_AT_ENTRY\"")
  .Attr("allow_soft_placement: bool = false")
  + .Attr("num_distributed_variables: int = 0")
  .Input("inputs: Tinputs")
  .Input("broadcast_inputs: Tbroadcast_inputs")
  .Input("variables: NumVariables * resource")
  