(*result)[id].push_back(ones_like);
        }
      }
    } else {
      (*result)[id].push_back(output_gradients[i]);
    }
  }
  return absl::OkStatus();
}

// TODO(agarwal): use an automatic mechanism for handling None arguments to
// gradient functions.
//
// Some gradient functions can accept None arguments for gradients. The

Marchesi, Max Melnick, Micael Carvalho, @mikowals, Mostafa Gazar, Nico Galoppo,
Nishant Agrawal, Petr Janda, Yuncheng Li, @raix852, Robert Rose,
@Robin-des-Bois, Rohit Girdhar, Sam Abrahams, satok16, Sergey Kishchenko, Sharkd
Tu, @shotat, Siddharth Agrawal, Simon Denel, @sono-bfio, SunYeop Lee, Thijs
Vogels, @tobegit3hub, @Undo1, Wang Yang, Wenjian Huang, Yaroslav Bulatov, Yuan