ReluModel, ReluGradModel, immediate_execution_ctx_.get(), {X.get()},
      UseFunction()));

  // Mathematically, Relu isn't differentiable at `0`. So `gradient_checker`
  // does not work with it.
  AbstractTensorHandlePtr Y;
  {
    AbstractTensorHandle* Y_raw;
    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
        immediate_execution_ctx_.get(), 0.0f, &Y_raw);

  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  EXPECT_EQ(dims[0], returned_dims[0]);
  EXPECT_EQ(dims[1], returned_dims[1]);

  // Try to set 'unknown' with unknown rank on the shape and see that
  // it doesn't change.
  TF_GraphSetTensorShape(graph, feed_out_0, /*dims=*/nullptr, -1, s);
  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  TF_GraphGetTensorShape(graph, feed_out_0, returned_dims, num_dims, s);