// TODO: b/310291615 - find a better way for platform support.
  arg_types.push_back(RankedTensorType::get({}, builder.getI32Type()));
  arg_locs.push_back(module_op.getLoc());
  for (const Value input_value : inputs) {
    arg_types.push_back(input_value.getType());
    arg_locs.push_back(input_value.getLoc());
  }

  // Identify result types.
  SmallVector<Type> result_types;

    return it->second;
  }

  int new_id = names_.size();
  names_.push_back(string(name));
  id_to_device_type_.push_back(std::make_unique<DeviceType>(""));
  DeviceType* device_type = id_to_device_type_.back().get();
  TF_RETURN_IF_ERROR(DeviceNameToDeviceType(names_.back(), device_type));

  is_cpu_.push_back(device_type->type_string() == DEVICE_CPU);
  is_gpu_.push_back(device_type->type_string() == DEVICE_GPU);

  TaggedValue valuef(3.f);
  TaggedValue valuei(int64_t(3));
  TaggedValue list = TaggedValue::List();
  TaggedValue tuple = TaggedValue::Tuple();
  tuple.tuple().push_back(TaggedValue(int64_t(310)));
  list.list().push_back(valuei);
  list.list().push_back(valuef);
  list.list().push_back(tuple);
  std::stringstream stream;
  stream << list;
  ASSERT_EQ(stream.str(), "[3, 3, (310, ), ]");
}

TEST(ValueTest, TestString) {

  *(config_1.mutable_tf_record()->mutable_path()) = "test_path_1";
  representative_dataset_configs.push_back(config_1);

  // Signature key is implicitly "serving_default".
  RepresentativeDatasetConfig config_2{};
  *(config_2.mutable_tf_record()->mutable_path()) = "test_path_2";
  representative_dataset_configs.push_back(config_2);

  const absl::StatusOr<
      absl::flat_hash_map<std::string, RepresentativeDatasetFile>>

      auto type = std::get<1>(it);
      if (value.getType() == type) {
        args.push_back(value);
      } else {
        if (IsCompatibleTypeWithTFLCastOp(value.getType()) &&
            IsCompatibleTypeWithTFLCastOp(type)) {
          auto cast = b.create<CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        } else {
          auto cast = b.create<TF::CastOp>(yield_op->getLoc(), type, value);

  auto shift = op.input(1);
  auto axis = op.input(2);
  auto grad_op = Roll(scope, grad_inputs[0], Neg(scope, shift), axis);
  grad_outputs->push_back(grad_op);
  grad_outputs->push_back(NoGradient());
  grad_outputs->push_back(NoGradient());
  return scope.status();
}
REGISTER_GRADIENT_OP("Roll", RollGrad);

}  // namespace
}  // namespace ops

};

TEST_F(PjRtExecutionUtilTest, PreparePjRtExecutableArguments) {
  std::vector<const Tensor*> inputs;
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {0, 0, 0}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {1, 2, 3}));
  inputs.push_back(CreateDeviceTensor<int32_t>(TensorShape({1, 3}), {4, 5, 6}));
  int num_missing_prefix_ctx_inputs = 2;
  std::vector<int> input_mapping{3, 4};

  for (int i = 0; i < num_inputs; i++) {
    func_inputs.push_back(op.input(i));
    input_dtypes.push_back(op.input_type(i));
  }

  func_inputs.insert(std::end(func_inputs), std::begin(grad_inputs),
                     std::end(grad_inputs));

  auto grad = SymbolicGradient(scope, func_inputs, input_dtypes, f);
  for (int i = 0; i < num_inputs; i++) {
    grad_outputs->push_back(grad[i]);
  }

  return scope.status();
}

          return WalkResult::interrupt();
        }

        if (llvm::find(func_metadata.inputs, value) ==
            func_metadata.inputs.end()) {
          func_metadata.inputs.push_back(value);
          func_metadata.input_devices.push_back(value_device);
        }
        continue;
      }

      Operation *defining_op = value.getDefiningOp();
      std::string defining_op_host = GetHost(defining_op);

    output_args_.push_back(arg);
  }
  // Parse the attributes.
  for (const OpDef::AttrDef& attr_def : op_def.attr()) {
    AttrSpec attr = AttrSpec::Create(attr_def);
    // Only non-inferred args are added as arguments.
    if (inferred_attrs.find(attr_def.name()) == inferred_attrs.end()) {
      argument_attrs_.push_back(attr);
    }
  }
}

}  // namespace generator