Xが4D、Yが2D、axisが1の場合に対応

sony · Apr 20, 2018 · 98e01e5 · 98e01e5
1 parent dfb35aa
commit 98e01e5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/python/test/function/test_broadcast_to.py b/python/test/function/test_broadcast_to.py
@@ -66,6 +66,12 @@ def ref_broadcast_to(x, y, axis):
                 t = y[np.newaxis, np.newaxis, :, np.newaxis]
                 t.transpose()
                 return np.broadcast_to(t, x.shape)
+        elif ys == 2:
+            if axis == 0:
+                pass
+            elif axis == 1:
+                t = y[np.newaxis, :, :, np.newaxis]
+                return np.broadcast_to(t, x.shape)
 
 
 
@@ -81,7 +87,7 @@ def ref_broadcast_to(x, y, axis):
     ((2, 3, 4, 5), (3), 1),
     ((2, 3, 4, 5), (4), 2),
     ((2, 3, 4, 5), (5), 3),
-    #((2, 3, 4, 5), (3, 4), 1),
+    ((2, 3, 4, 5), (3, 4), 1),
     ((2, 3, 4, 5), (4, 5), 2),
     #((2, 3, 4, 5), (2, 3, 4), 0),
     ((2, 3, 4, 5), (3, 4, 5), 1),

diff --git a/src/nbla/function/broadcast_to.cpp b/src/nbla/function/broadcast_to.cpp
@@ -255,7 +255,17 @@ void BroadcastTo<T>::forward_impl(const Variables &inputs, const Variables &outp
 							// X: (2,3,4,5) Y: (2,3) axis=0
 							break;
 						case 1:
-							// X: (2,3,4,5) Y: (3,4) axis=1
+							{
+								// X: (2,3,4,5) Y: (3,4) axis=1
+								const Size_t block_size = xs[2]*xs[3];
+								const Size_t slice_size = xs[1]*block_size;
+								for (Size_t i=0; i<xs[0]; i++) {
+									copy_buf_vertically_to_block(
+											&z[i*slice_size], y,
+											xs[1], ys[1],
+											xs[2], xs[3]);
+								}
+							}
 							break;
 						case 2:
 							// X: (2,3,4,5) Y: (4,5) axis=2