AbstractTensorHandle* expand_dims_outputs;
  TF_RETURN_IF_ERROR(
      ops::ExpandDims(ctx, vec, dim.get(), &expand_dims_outputs, "ExpandDims"));
  TF_RETURN_IF_ERROR(
      ops::Mul(ctx, expand_dims_outputs, mat, &outputs[0], "Mul"));
  expand_dims_outputs->Unref();
  return absl::OkStatus();
}

class SparseSoftmaxCrossEntropyWithLogitsGradientFunction
    : public GradientFunction {
 public:

    /* Given upstream grad U and a mul op A*B, the gradients are:
     *
     *    dA = U * B
     *    dB = A * U
     *
     */

    AbstractTensorHandle* upstream_grad = grad_outputs[0];

    // Gradient for A
    std::string name = "Mul_Grad_A";
    TF_RETURN_IF_ERROR(Mul(ctx, upstream_grad, forward_inputs_[1],

  EXPECT_EQ(2, product[0]);
  EXPECT_EQ(4, product[1]);
  EXPECT_EQ(6, product[2]);
  EXPECT_EQ(8, product[3]);

  // When we turn on grappler, confirm that the tf.Add has been rewritten into a
  // tf.Mul.
  // This capability of checking the executed op names is currently only enabled
  // for TFRT debug build, for performance and simplicity reasons.
  if (use_tfrt) {
    TF_Buffer* buf = TF_NewBuffer();

TEST_F(CApiGradientsTest, MultipleCallsToAddGradients) {
  const float X = 3.0f, Y = 7.0f;
  TF_Operation* x = Placeholder(graph_, s_, "x", TF_FLOAT);
  TF_Operation* y = Placeholder(graph_, s_, "y", TF_FLOAT);
  TF_Operation* xy = Mul(x, y, graph_, s_, "xy");
  TF_Output dxy_dx, dxy_dy;

  TF_Output outputs[1] = {{xy, 0}};
  TF_Output inputs[1] = {{x, 0}};
  TF_AddGradients(graph_, outputs, 1, inputs, 1, nullptr, s_, &dxy_dx);

*   `tf.data`:

    *   Moved option `warm_start` from `tf.data.experimental.OptimizationOptions` to `tf.data.Options`.

*   `tf.lite`:

    *   `sub_op` and `mul_op` support broadcasting up to 6 dimensions.