Use explicit types in test

Author: trevor-m

xla/pjrt/gpu/se_gpu_pjrt_client_test.cc

@@ -1713,41 +1713,45 @@ TEST(StreamExecutorGpuClientTest, NvshmemMemoryTest) {
   client_options.allowed_devices = {0};
   client_options.num_nodes = 1;
   client_options.kv_store = std::make_shared<InMemoryKeyValueStore>();
-  TF_ASSERT_OK_AND_ASSIGN(auto client,
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
                           GetStreamExecutorGpuClient(client_options));
   xla::CompileOptions options;
   options.executable_build_options.mutable_debug_options()
       ->set_xla_gpu_experimental_enable_nvshmem(true);
 
-  TF_ASSERT_OK_AND_ASSIGN(auto executable,
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::PjRtLoadedExecutable> executable,
                           CompileExecutable(kProgram, *client, options));
   std::vector<int32_t> data{1, 2, 3, 4};
   Shape shape = ShapeUtil::MakeShapeWithDenseLayout(S32, {1, 4},
                                                     /*minor_to_major=*/{1, 0});
   shape.mutable_layout()->set_memory_space(Layout::kDefaultMemorySpace);
 
-  auto device = client->addressable_devices()[0];
+  PjRtDevice* const device = client->addressable_devices()[0];
   TF_EXPECT_OK(device->default_memory_space());
   TF_ASSERT_OK_AND_ASSIGN(
-      auto input, client->BufferFromHostBuffer(
-                      data.data(), shape.element_type(), shape.dimensions(),
-                      /*byte_strides=*/std::nullopt,
-                      PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
-                      /*on_done_with_host_buffer=*/nullptr, device));
+      std::unique_ptr<PjRtBuffer> input,
+      client->BufferFromHostBuffer(
+          data.data(), shape.element_type(), shape.dimensions(),
+          /*byte_strides=*/std::nullopt,
+          PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
+          /*on_done_with_host_buffer=*/nullptr, *device->default_memory_space(),
+          /*device_layout=*/nullptr));
   EXPECT_EQ(input->memory_space()->kind(), "device");
 
-  TF_ASSERT_OK_AND_ASSIGN(auto memory_kinds,
-                          executable->GetOutputMemoryKinds());
+  TF_ASSERT_OK_AND_ASSIGN(
+      std::vector<std::vector<absl::string_view>> memory_kinds,
+      executable->GetOutputMemoryKinds());
   EXPECT_EQ(memory_kinds.size(), 1);
   EXPECT_EQ(memory_kinds[0].size(), 1);
   EXPECT_EQ(memory_kinds[0][0], "device");
 
   TF_ASSERT_OK_AND_ASSIGN(
-      auto result, executable->Execute({{input.get()}}, ExecuteOptions()));
+      std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> result,
+      executable->Execute({{input.get()}}, ExecuteOptions()));
   std::vector<std::unique_ptr<xla::PjRtBuffer>>& result_buffers = result[0];
   EXPECT_EQ(result_buffers[0]->memory_space()->kind(), "device");
   Shape result_shape = result_buffers[0]->on_device_shape();
-  auto memory_space = result_shape.layout().memory_space();
+  int64_t memory_space = result_shape.layout().memory_space();
   EXPECT_EQ(memory_space, 1);
 }