support axis argument in nn.glu (#2879)

* support axis argument in nn.glu

* also add basic correctness test

* Update nn_test.py

jax/nn/functions.py

@@ -189,7 +189,8 @@ def glu(x, axis=-1):
   """Gated linear unit activation function."""
   size = x.shape[axis]
   assert size % 2 == 0, "axis size must be divisible by 2"
-  return x[..., :size // 2] * sigmoid(x[..., size // 2:])
+  x1, x2 = np.split(x, 2, axis)
+  return x1 * sigmoid(x2)
 
 # other functions
 

tests/nn_test.py

@@ -88,6 +88,10 @@ def testEluGrad(self):
   def testEluValue(self):
     val = nn.elu(1e4)
     self.assertAllClose(val, 1e4, check_dtypes=False)
+
+  def testGluValue(self):
+    val = nn.glu(np.array([1.0, 0.0]))
+    self.assertAllClose(val, np.array([0.5]), check_dtypes=True)
 
   @parameterized.parameters(*itertools.product(
       (np.float32, np.bfloat16, np.float16),