From 96cca6c5b90788a70356b65f24ab6beb952f1f9e Mon Sep 17 00:00:00 2001
From: Wenzel Jakob <wenzel.jakob@epfl.ch>
Date: Fri, 20 Sep 2024 10:40:25 +0900
Subject: [PATCH] follow strict tensorflow alignment requirements

---
 tests/test_ndarray.cpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/test_ndarray.cpp b/tests/test_ndarray.cpp
index 42b9e0d1..8c479aa2 100644
--- a/tests/test_ndarray.cpp
+++ b/tests/test_ndarray.cpp
@@ -312,15 +312,18 @@ NB_MODULE(test_ndarray_ext, m) {
     });
 
     m.def("ret_tensorflow", []() {
-        float *f = new float[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
+        struct alignas(256) Buf {
+            float f[8];
+        };
+        Buf *buf = new Buf({ 1, 2, 3, 4, 5, 6, 7, 8 });
         size_t shape[2] = { 2, 4 };
 
-        nb::capsule deleter(f, [](void *data) noexcept {
+        nb::capsule deleter(buf, [](void *data) noexcept {
            destruct_count++;
-           delete[] (float *) data;
+           delete[] (Buf *) data;
         });
 
-        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(f, 2, shape,
+        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(buf->f, 2, shape,
                                                                    deleter);
     });