double abs_error) {
  auto num_inputs = inputs.size();
  std::vector<AbstractTensorHandle*> outputs(num_inputs);
  auto s = RunModel(grad_model, ctx, inputs, absl::MakeSpan(outputs),
                    /*use_function=*/use_function);
  ASSERT_EQ(errors::OK, s.code()) << s.message();

  for (int i = 0; i < num_inputs; ++i) {
    if (!outputs[i]) continue;

  const Node* identity_node = node_name_index.at("identity_op");
  EXPECT_EQ(identity_node->op_def().name(), "Identity");
  EXPECT_EQ(identity_node->attrs().FindByString("T")->type(), DT_FLOAT);

  EXPECT_EQ(identity_node->num_inputs(), 1);
  const Node* identity_input_node = nullptr;
  TF_EXPECT_OK(identity_node->input_node(0, &identity_input_node));
  EXPECT_EQ(identity_input_node->name(), "_arg0");

  using tensorflow::shape_inference::ShapeHandle;

  const int num_inputs = input_shapes->num_items;
  NodeDef node_def;
  tensorflow::ImmediateExecutionOperation* op = tensorflow::unwrap(tfe_op);
  node_def.set_name(op->Name());
  node_def.set_op(op->Name());
  for (int i = 0; i < num_inputs; ++i) {
    node_def.add_input("dummy_input");
  }

    if (data_type.getIntOrFloatBitWidth() == 64) continue;

    const auto& block_arg = replicate.GetBody().getArgument(replicate_arg);

    int64_t num_inputs = 0;
    if (replicate.IsReplicatedBlockArgument(block_arg)) {
      num_inputs = num_replicas;
    } else {
      num_inputs = 1;
    }

    // We have found a mirrored variable which is an input to the replicated

namespace tensorflow {

namespace {

class WhileLoopTest : public ::testing::Test {
 protected:
  WhileLoopTest() : scope_(Scope::NewRootScope()) {}

  void Init(int num_inputs, DataType dtype = DT_INT32) {
    for (int i = 0; i < num_inputs; ++i) {
      inputs_.push_back(ops::Placeholder(scope_, dtype));
    }
  }

  void CreateLoop(const ops::CondGraphBuilderFn& cond,

                  const std::vector<Output>& grad_inputs,
                  std::vector<Output>* grad_outputs) {
  if (op.num_inputs() < 3) {
    return errors::InvalidArgument("SplitV requires 3 arguments");
  }
  grad_outputs->push_back(Concat(scope, grad_inputs, op.input(2)));
  for (int i = 0; i < op.num_inputs() - 1; ++i) {
    grad_outputs->push_back(NoGradient());
  }
  return scope.status();
}

  return cc_ctx->num_inputs();
}

void TF_ShapeInferenceContextGetInput(TF_ShapeInferenceContext* ctx, int i,
                                      TF_ShapeHandle* handle,
                                      TF_Status* status) {
  TF_SetStatus(status, TF_OK, "");
  auto* cc_ctx = reinterpret_cast<InferenceContext*>(ctx);
  if (i < 0 || i >= cc_ctx->num_inputs()) {

    std::vector<const Edge*> in_edges;
    TF_RETURN_IF_ERROR(launch->input_edges(&in_edges));

    const int num_inputs = in_edges.size();
    const int variable_start_index = xla_function_info.variable_start_index;
    const int num_variables = num_inputs - variable_start_index;
    const int num_args = variable_start_index;

    VLOG(4) << "Launch node '" << launch->name() << "'"

      : TensorHandleConvertible(
            testing::CreateTensorHandle(ctx, DT_FLOAT, {2, 4}, value)) {}
};

FunctionDef FuncDefWithNumInputsOutputs(int num_inputs, int num_outputs) {
  FunctionDef func;
  OpDef* signature = func.mutable_signature();
  for (int i = 0; i < num_inputs; ++i) {
    signature->add_input_arg();
  }
  for (int i = 0; i < num_outputs; ++i) {
    signature->add_output_arg();
  }
  return func;
}

  return variable_lookup;
}

}  // anonymous namespace

std::vector<const Tensor*> InputsFromContext(OpKernelContext* ctx) {
  std::vector<const Tensor*> inputs;
  inputs.reserve(ctx->num_inputs());
  for (int input_idx = 0; input_idx < ctx->num_inputs(); input_idx++) {
    inputs.push_back(&ctx->input(input_idx));
  }
  return inputs;
}

absl::StatusOr<std::vector<int>> GetConstantInputIndicesFromContext(