message: Allow broadcasting of rank 0 tensors
  Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep
  Maniphest Tasks: T64481
diff: @@ -22,8 +22,12 @@ REGISTER_OP("PopdistBroadcast")
  .Output("sum: T")
  .SetShapeFn([](shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle output;
  +
  + if (c->Rank(c->input(0)) > 0) {
  TF_RETURN_IF_ERROR(
  c->ReplaceDim(c->input(0), 0, c->UnknownDim(), &output));
  + }
  +
  c->set_output(0, output);
  return Status::OK();
  });
  