ctx_.get(), B_vals, B_dims, 2, &B_raw);
    ASSERT_EQ(errors::OK, s.code()) << s.message();
    B.reset(B_raw);
  }

  float expected_dA[4] = {-.5f, 2.0f, -.5f, 2.0f};
  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndManualGradients(
      MatMulModel, ctx_.get(), {A.get(), B.get()}, 0, expected_dA, 4,
      UseFunction()));
}

TEST_P(GradientCheckerTest, TestMul) {
  AbstractTensorHandlePtr x;