Merge pull request #2701 from cgarciae:cleanup-dense-general

PiperOrigin-RevId: 495114816

Author: Flax Authors

flax/linen/linear.py

@@ -14,7 +14,6 @@
 
 """Linear modules."""
 
-import abc
 import dataclasses
 from typing import (Any, Callable, Iterable, List, Optional, Sequence, Tuple,
                     Union)
@@ -31,6 +30,7 @@
 from jax import ShapedArray
 import jax.numpy as jnp
 import numpy as np
+import jax
 
 
 PRNGKey = Any
@@ -107,12 +107,11 @@ def __call__(self, inputs: Array) -> Array:
     n_axis, n_features = len(axis), len(features)
 
     def kernel_init_wrap(rng, shape, dtype=jnp.float32):
-      size_batch_dims = np.prod(shape[:n_batch_dims], dtype=np.int32)
-      flat_shape = (np.prod(shape[n_batch_dims:n_axis + n_batch_dims]),
+      flat_shape = (np.prod(shape[:n_batch_dims]) *
+                    np.prod(shape[n_batch_dims:n_axis + n_batch_dims]),
                     np.prod(shape[-n_features:]),)
-      kernel = jnp.concatenate(
-          [self.kernel_init(rng, flat_shape, dtype)
-           for rng in random.split(rng, size_batch_dims)], axis=0)
+      flat_shape = jax.tree_map(int, flat_shape)
+      kernel = self.kernel_init(rng, flat_shape, dtype)
       return jnp.reshape(kernel, shape)
 
     batch_shape = tuple(inputs.shape[ax] for ax in batch_dims)
@@ -129,11 +128,10 @@ def kernel_init_wrap(rng, shape, dtype=jnp.float32):
 
     if self.use_bias:
       def bias_init_wrap(rng, shape, dtype=jnp.float32):
-        size_batch_dims = np.prod(shape[:n_batch_dims], dtype=np.int32)
-        flat_shape = (np.prod(shape[-n_features:]),)
-        bias = jnp.concatenate(
-            [self.bias_init(rng, flat_shape, dtype)
-             for rng in random.split(rng, size_batch_dims)], axis=0)
+        flat_shape = (np.prod(shape[:n_batch_dims]) *
+                      np.prod(shape[-n_features:]),)
+        flat_shape = jax.tree_map(int, flat_shape)
+        bias = self.bias_init(rng, flat_shape, dtype)
         return jnp.reshape(bias, shape)
 
       bias = self.param('bias', bias_init_wrap, batch_shape + features,

tests/linen/linen_linear_test.py

@@ -77,13 +77,13 @@ def test_dense_is_dense_general(self):
         use_bias=True,
         bias_init=initializers.normal(),
     )
-    y1, params = dense_module.init_with_output(dict(params=random.PRNGKey(1)), x)
+    y1, _ = dense_module.init_with_output(dict(params=random.PRNGKey(1)), x)
     dg_module = nn.DenseGeneral(
         features=4,
         use_bias=True,
         bias_init=initializers.normal(),
     )
-    y2 = dg_module.apply(params, x)
+    y2, _ = dg_module.init_with_output(dict(params=random.PRNGKey(1)), x)
 
     np.testing.assert_allclose(y1, y2)
 
@@ -141,8 +141,7 @@ def _counter_init(rng, shape, dtype, state):
         kernel_init=counter_init,
     )
     y, _ = dg_module.init_with_output(rng, x)
-    target = np.concatenate(
-        [np.full((1, 1, 7), 16.), np.full((1, 1, 7), 31.)], axis=0)
+    target = np.full((2, 1, 7), 16.)
     np.testing.assert_allclose(y, target)
 
   @parameterized.parameters([((-2, 3), (), 'bijk,jklm->bilm'),