Skip to content

Commit f681878

Browse files
committed
[ET-VK][5/n] make_seq_tensor in codegen
increasing sequence is very useful for development, particularly for "slicing" and "indexing" operations. Differential Revision: [D56095314](https://our.internmc.facebook.com/intern/diff/D56095314/) ghstack-source-id: 222417564 Pull Request resolved: #3088
1 parent 029225f commit f681878

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

backends/vulkan/test/op_tests/utils/codegen_base.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,25 @@ def generate_suite_cpp(self) -> str:
218218
return at::rand(sizes, at::device(at::kCPU).dtype(dtype)) * (high - low) + low;
219219
}}
220220
221+
222+
at::Tensor make_seq_tensor(
223+
std::vector<int64_t> sizes,
224+
at::ScalarType dtype = at::kFloat) {{
225+
int64_t n = 1;
226+
for (auto size: sizes) {{
227+
n *= size;
228+
}}
229+
230+
std::vector<float> values(n);
231+
for (int i=0;i<n;i++) {{
232+
values[i] = (float) i;
233+
}}
234+
235+
// from_blob doesn't take ownership of data. Hence must create a copy as
236+
// "values" will go out of scope.
237+
return at::from_blob(values.data(), sizes, dtype).detach().clone();
238+
}}
239+
221240
{test_suites_cpp}
222241
"""
223242

0 commit comments

Comments
 (0)