xref: /aosp_15_r20/external/pytorch/aten/src/ATen/test/lazy_tensor_test.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 
3 #include <ATen/ATen.h>
4 
LazyTensorTest(c10::DispatchKey dispatch_key,at::DeviceType device_type)5 void LazyTensorTest(c10::DispatchKey dispatch_key, at::DeviceType device_type) {
6   auto tensor_impl =
7       c10::make_intrusive<c10::TensorImpl, c10::UndefinedTensorImpl>(
8           dispatch_key,
9           caffe2::TypeMeta::Make<float>(),
10           at::Device(device_type, 0));
11   at::Tensor t(std::move(tensor_impl));
12   ASSERT_TRUE(t.device() == at::Device(device_type, 0));
13 }
14 
15 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
TEST(XlaTensorTest,TestNoStorage)16 TEST(XlaTensorTest, TestNoStorage) {
17   LazyTensorTest(at::DispatchKey::XLA, at::DeviceType::XLA);
18 }
19 
20 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
TEST(LazyTensorTest,TestNoStorage)21 TEST(LazyTensorTest, TestNoStorage) {
22   LazyTensorTest(at::DispatchKey::Lazy, at::DeviceType::Lazy);
23 }
24