Repository: chainer/chainer

Issue:
`_chainerx_is_supported` in `F.batch_normalization` is wrong
It should test `x.ndim`.

* Conditions
```
Platform: Linux-4.4.0-103-generic-x86_64-with-debian-stretch-sid
Chainer: 7.0.0rc1
ChainerX: Available
NumPy: 1.17.2
CuPy:
  CuPy Version          : 7.0.0rc1
  CUDA Root             : /usr/local/cuda
  CUDA Build Version    : 9000
  CUDA Driver Version   : 9000
  CUDA Runtime Version  : 9000
  cuDNN Build Version   : 7300
  cuDNN Version         : 7300
  NCCL Build Version    : 2005
  NCCL Runtime Version  : (unknown)
iDeep: 2.0.0.post3
```

* Code to reproduce
```python
import chainer
import chainerx
import numpy as np


np.random.seed(0)
x = np.random.randn(8, 9, 2, 3, 4).astype('f')
devices = ['@cupy:0', 'cuda:0']
for axis in [
        (0,),  # fallback
        (0, 2, 3, 4),  # cudnn
        (0, 2, 3),  # bad
]:
    print(axis)
    ys = [
        chainer.get_device('@numpy').send(
            chainer.links.BatchNormalization(axis=axis).to_device(device)(
                chainer.get_device(device).send(x)
            ).array)
        for device in devices
    ]
    np.testing.assert_allclose(ys[0], ys[1], rtol=1e-4, atol=1e-4, verbose=False)
```

* Error messages, stack traces, or logs
```
(0,)
(0, 2, 3, 4)
(0, 2, 3)
Traceback (most recent call last):
  ...
AssertionError:
Not equal to tolerance rtol=0.0001, atol=0.0001

Mismatch: 100%
Max absolute difference: 0.6870811
Max relative difference: 293.6555
```


Discussion:

Resolving patch:
diff --git a/chainer/functions/normalization/batch_normalization.py b/chainer/functions/normalization/batch_normalization.py
--- a/chainer/functions/normalization/batch_normalization.py
+++ b/chainer/functions/normalization/batch_normalization.py
@@ -355,7 +355,7 @@ def forward_chainerx(self, inputs):
 
         x, gamma, beta = inputs
         axis_chx = _chainerx_compute_axis(x.ndim, gamma.ndim, self.axis)
-        if not _chainerx_is_supported(x.device, axis_chx):
+        if not _chainerx_is_supported(x, axis_chx):
             return chainer.Fallback
 
         y = chainerx.batch_norm(
@@ -578,7 +578,7 @@ def forward_chainerx(self, inputs):
 
         x, gamma, beta, mean, var = inputs
         axis_chx = _chainerx_compute_axis(x.ndim, gamma.ndim, self.axis)
-        if not _chainerx_is_supported(x.device, axis_chx):
+        if not _chainerx_is_supported(x, axis_chx):
             return chainer.Fallback
 
         y = chainerx.fixed_batch_norm(
@@ -797,20 +797,19 @@ def _chainerx_compute_axis(x_ndim, gamma_ndim, axis):
         else axis if isinstance(axis, tuple)
         else (axis,))
     axis_chx = _compute_axis(x_ndim, gamma_ndim, axis_chx)
+    assert isinstance(axis_chx, tuple)
     return axis_chx
 
 
-def _chainerx_is_supported(device, axis_chx):
+def _chainerx_is_supported(x, axis_chx):
     # Checks if the input configuration is supported in ChainerX
-    axis_ndim_chx = len(axis_chx)
+    device = x.device
     if device.backend.name == 'cuda':
         # cuDNN batch norm restriction
-        if not ((axis_ndim_chx == 3 and axis_chx[0] == 0
-                 and axis_chx[1] == 2 and axis_chx[2] == 3)
-                or (axis_ndim_chx == 4 and axis_chx[0] == 0
-                    and axis_chx[1] == 2 and axis_chx[2] == 3
-                    and axis_chx[3] == 4)):
-            return False
+        return (x.ndim, axis_chx) in [
+            (4, (0, 2, 3)),
+            (5, (0, 2, 3, 4)),
+        ]
     return True
 
 


Test patch: