Repository: google/jax

Issue:
Need a function for concat layers in stax, maybe FanInCat
Sometimes we need concat many layers. So like parallel, FanOut and FanInSum, we need FanInCat


Discussion:

Resolving patch:
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -212,6 +212,18 @@ def FanInSum():
 FanInSum = FanInSum()
 
 
+def FanInConcat(axis=-1):
+  """Layer construction function for a fan-in concatenation layer."""
+  def init_fun(input_shape):
+    ax = axis % len(input_shape[0])
+    concat_size = sum(shape[ax] for shape in input_shape)
+    out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
+    return out_shape, ()
+  def apply_fun(params, inputs, rng=None):
+    return np.concatenate(inputs, axis)
+  return init_fun, apply_fun
+
+
 def Dropout(rate, mode='train'):
   """Layer construction function for a dropout layer with given rate."""
   def init_fun(input_shape):
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 
 setup(
     name='jax',
-    version='0.1.14',
+    version='0.1.15',
     description='Differentiate, compile, and transform Numpy code.',
     author='JAX team',
     author_email='jax-dev@google.com',


Test patch:
diff --git a/tests/stax_test.py b/tests/stax_test.py
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -30,9 +30,18 @@
 config.parse_flags_with_absl()
 
 
+def random_inputs(rng, input_shape):
+  if type(input_shape) is tuple:
+    return rng.randn(*input_shape).astype(onp.float32)
+  elif type(input_shape) is list:
+    return [random_inputs(rng, shape) for shape in input_shape]
+  else:
+    raise TypeError(type(input_shape))
+
+
 def _CheckShapeAgreement(test_case, init_fun, apply_fun, input_shape):
   result_shape, params = init_fun(input_shape)
-  inputs = onp.random.RandomState(0).randn(*input_shape).astype("float32")
+  inputs = random_inputs(onp.random.RandomState(0), input_shape)
   rng_key = random.PRNGKey(0)
   result = apply_fun(params, inputs, rng_key)
   test_case.assertEqual(result.shape, result_shape)
@@ -130,6 +139,26 @@ def testDropoutShape(self, input_shape):
     init_fun, apply_fun = stax.Dropout(0.9)
     _CheckShapeAgreement(self, init_fun, apply_fun, input_shape)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_input_shape={}".format(input_shape),
+       "input_shape": input_shape}
+      for input_shape in [(3, 4), (2, 5, 6, 1)]))
+  def testFanInSum(self, input_shape):
+    init_fun, apply_fun = stax.FanInSum
+    _CheckShapeAgreement(self, init_fun, apply_fun, [input_shape, input_shape])
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_inshapes={}_axis={}".format(input_shapes, axis),
+       "input_shapes": input_shapes, "axis": axis}
+      for input_shapes, axis in [
+          ([(2, 3), (2, 1)], 1),
+          ([(2, 3), (2, 1)], -1),
+          ([(1, 2, 4), (1, 1, 4)], 1),
+      ]))
+  def testFanInConcat(self, input_shapes, axis):
+    init_fun, apply_fun = stax.FanInConcat(axis)
+    _CheckShapeAgreement(self, init_fun, apply_fun, input_shapes)
+
 
 if __name__ == "__main__":
   absltest.main()