From a478ccee068d5443b0bbbeb810f2a808a9f1e068 Mon Sep 17 00:00:00 2001
From: He Sichao <1310722434@qq.com>
Date: Sun, 3 Mar 2024 12:38:01 +0800
Subject: [PATCH] Update

---
 brainpy/_src/dnn/conv.py                      | 11 ++++++++++-
 brainpy/_src/dnn/tests/test_activation.py     |  2 +-
 brainpy/_src/dnn/tests/test_conv_layers.py    | 11 ++++++-----
 brainpy/_src/dnn/tests/test_function.py       |  6 ++----
 brainpy/_src/dnn/tests/test_normalization.py  |  3 ++-
 brainpy/_src/dnn/tests/test_pooling_layers.py |  2 +-
 6 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/brainpy/_src/dnn/conv.py b/brainpy/_src/dnn/conv.py
index deead1f3b..e4b6e25d2 100644
--- a/brainpy/_src/dnn/conv.py
+++ b/brainpy/_src/dnn/conv.py
@@ -160,7 +160,7 @@ def update(self, x):
     nonbatching = False
     if x.ndim == self.num_spatial_dims + 1:
       nonbatching = True
-      x = x.unsqueeze(0)
+      x = bm.unsqueeze(x, 0)
     w = self.w.value
     if self.mask is not None:
       try:
@@ -190,6 +190,9 @@ def __repr__(self):
 class Conv1d(_GeneralConv):
   """One-dimensional convolution.
 
+  The input should a 2d array with the shape of ``[H, C]``, or
+  a 3d array with the shape of ``[B, H, C]``, where ``H`` is the feature size.
+
   Parameters
   ----------
   in_channels: int
@@ -282,6 +285,9 @@ def _check_input_dim(self, x):
 class Conv2d(_GeneralConv):
   """Two-dimensional convolution.
 
+  The input should a 3d array with the shape of ``[H, W, C]``, or
+  a 4d array with the shape of ``[B, H, W, C]``.
+
   Parameters
   ----------
   in_channels: int
@@ -375,6 +381,9 @@ def _check_input_dim(self, x):
 class Conv3d(_GeneralConv):
   """Three-dimensional convolution.
 
+  The input should a 3d array with the shape of ``[H, W, D, C]``, or
+  a 4d array with the shape of ``[B, H, W, D, C]``.
+
   Parameters
   ----------
   in_channels: int
diff --git a/brainpy/_src/dnn/tests/test_activation.py b/brainpy/_src/dnn/tests/test_activation.py
index ba2a49efd..7a0fa57af 100644
--- a/brainpy/_src/dnn/tests/test_activation.py
+++ b/brainpy/_src/dnn/tests/test_activation.py
@@ -1,5 +1,5 @@
-from absl.testing import parameterized
 from absl.testing import absltest
+from absl.testing import parameterized
 import brainpy as bp
 import brainpy.math as bm
 
diff --git a/brainpy/_src/dnn/tests/test_conv_layers.py b/brainpy/_src/dnn/tests/test_conv_layers.py
index 3c9fdfa87..05f523622 100644
--- a/brainpy/_src/dnn/tests/test_conv_layers.py
+++ b/brainpy/_src/dnn/tests/test_conv_layers.py
@@ -1,17 +1,15 @@
 # -*- coding: utf-8 -*-
 
-from unittest import TestCase
-from absl.testing import absltest
 import jax.numpy as jnp
-import brainpy.math as bm
+from absl.testing import absltest
 from absl.testing import parameterized
+
 import brainpy as bp
 import brainpy.math as bm
 
 
 class TestConv(parameterized.TestCase):
   def test_Conv2D_img(self):
-    bm.random.seed()
     img = jnp.zeros((2, 200, 198, 4))
     for k in range(4):
       x = 30 + 60 * k
@@ -24,6 +22,7 @@ def test_Conv2D_img(self):
                              strides=(2, 1), padding='VALID', groups=4)
       out = net(img)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 99, 196, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(img)[0, :, :, 0])
@@ -31,7 +30,6 @@ def test_Conv2D_img(self):
     bm.clear_buffer_memory()
 
   def test_conv1D(self):
-    bm.random.seed()
     with bp.math.training_environment():
       model = bp.layers.Conv1d(in_channels=3, out_channels=32, kernel_size=(3,))
 
@@ -39,6 +37,7 @@ def test_conv1D(self):
 
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(out)[0, :, :])
@@ -54,6 +53,7 @@ def test_conv2D(self):
 
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 5, 32))
       # print("First output channel:")
       # plt.figure(figsize=(10, 10))
       # plt.imshow(np.array(out)[0, :, :, 31])
@@ -67,6 +67,7 @@ def test_conv3D(self):
       input = bp.math.ones((2, 5, 5, 5, 3))
       out = model(input)
       print("out shape: ", out.shape)
+      self.assertEqual(out.shape, (2, 5, 5, 5, 32))
     bm.clear_buffer_memory()
 
 
diff --git a/brainpy/_src/dnn/tests/test_function.py b/brainpy/_src/dnn/tests/test_function.py
index 269fec441..9ad15938d 100644
--- a/brainpy/_src/dnn/tests/test_function.py
+++ b/brainpy/_src/dnn/tests/test_function.py
@@ -1,12 +1,10 @@
 # -*- coding: utf-8 -*-
 
-from unittest import TestCase
-
-import jax.numpy as jnp
-import brainpy.math as bm
 from absl.testing import absltest
 from absl.testing import parameterized
+
 import brainpy as bp
+import brainpy.math as bm
 
 
 class TestFunction(parameterized.TestCase):
diff --git a/brainpy/_src/dnn/tests/test_normalization.py b/brainpy/_src/dnn/tests/test_normalization.py
index fdc5b34e3..e76b3616b 100644
--- a/brainpy/_src/dnn/tests/test_normalization.py
+++ b/brainpy/_src/dnn/tests/test_normalization.py
@@ -1,7 +1,8 @@
-import brainpy.math as bm
 from absl.testing import parameterized
 from absl.testing import absltest
+
 import brainpy as bp
+import brainpy.math as bm
 
 
 class Test_Normalization(parameterized.TestCase):
diff --git a/brainpy/_src/dnn/tests/test_pooling_layers.py b/brainpy/_src/dnn/tests/test_pooling_layers.py
index 34f8f5cd5..5748edd8b 100644
--- a/brainpy/_src/dnn/tests/test_pooling_layers.py
+++ b/brainpy/_src/dnn/tests/test_pooling_layers.py
@@ -3,8 +3,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from absl.testing import parameterized
 from absl.testing import absltest
+from absl.testing import parameterized
 
 import brainpy as bp
 import brainpy.math as bm