How to remove units from a hidden layer without ScopeParamShapeError?

[samuela]

I'd like to take a large network and select a smaller sub-network out of it by only selecting to keep certain units. Is there a way to accomplish such a thing in flax while avoiding ScopeParamShapeError?

Here's a quick code example that builds a small model and then attempts to remove a single unit from the hidden layer:

import flax.linen as nn
import jax.numpy as jnp
from flax import traverse_util
from flax.core import freeze, unfreeze
from jax import random
from jax._src.tree_util import tree_map

def flatten_params(params):
  return {"/".join(k): v for k, v in traverse_util.flatten_dict(unfreeze(params)).items()}

def unflatten_params(flat_params):
  return freeze(
      traverse_util.unflatten_dict({tuple(k.split("/")): v
                                    for k, v in flat_params.items()}))

class _net(nn.Module):
  @nn.compact
  def __call__(self, x):
    # A straightforward one hidden layer network...
    x = nn.Dense(13)(x)
    x = nn.relu(x)
    x = nn.Dense(5)(x)
    x = nn.log_softmax(x)
    return x

net = _net()

# Get the initial parameters...
init_params = net.init(random.PRNGKey(0), jnp.zeros((1, 7)))
init_params_flat = flatten_params(init_params)

# Construct a sub-network by dropping the last unit of the first hidden layer...
sub_params = unflatten_params({
    "params/Dense_0/bias":
    init_params_flat["params/Dense_0/bias"][:12],
    "params/Dense_0/kernel":
    init_params_flat["params/Dense_0/kernel"][:, :12],
    "params/Dense_1/bias":
    init_params_flat["params/Dense_1/bias"],
    "params/Dense_1/kernel":
    init_params_flat["params/Dense_1/kernel"][:12, :],
})

print(tree_map(jnp.shape, init_params_flat))
print(tree_map(jnp.shape, flatten_params(sub_params)))

# This works:
net.apply(init_params, jnp.zeros((1, 7)))

# This doesn't:
net.apply(sub_params, jnp.zeros((1, 7)))

All of the shapes of sub_params arrays check out: mathematically they should compute the subnetwork just as desired. However flax gives me an error:

{'params/Dense_0/bias': (13,), 'params/Dense_0/kernel': (7, 13), 'params/Dense_1/bias': (5,), 'params/Dense_1/kernel': (13, 5)}
{'params/Dense_0/bias': (12,), 'params/Dense_0/kernel': (7, 12), 'params/Dense_1/bias': (5,), 'params/Dense_1/kernel': (12, 5)}
---------------------------------------------------------------------------
ScopeParamShapeError                      Traceback (most recent call last)
~/dev/research/research/lottery/flax_bug.py in <module>
     49 
     50 # This doesn't:
---> 51 net.apply(sub_params, jnp.zeros((1, 7)))

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/module.py in apply(self, variables, rngs, method, mutable, capture_intermediates, *args, **kwargs)
    934       method = self.__call__
    935     method = _get_unbound_fn(method)
--> 936     return apply(
    937         method, self,
    938         mutable=mutable, capture_intermediates=capture_intermediates

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/core/scope.py in wrapper(variables, rngs, *args, **kwargs)
    685               **kwargs) -> Union[Any, Tuple[Any, VariableDict]]:
    686     with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
--> 687       y = fn(root, *args, **kwargs)
    688     if mutable is not False:
    689       return y, root.mutable_variables()

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/module.py in scope_fn(scope, *args, **kwargs)
   1176     _context.capture_stack.append(capture_intermediates)
   1177     try:
-> 1178       return fn(module.clone(parent=scope), *args, **kwargs)
   1179     finally:
   1180       _context.capture_stack.pop()

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

~/dev/research/research/lottery/flax_bug.py in __call__(self, x)
     18   def __call__(self, x):
     19     # A straightforward one hidden layer network...
---> 20     x = nn.Dense(13)(x)
     21     x = nn.relu(x)
     22     x = nn.Dense(5)(x)

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/linear.py in __call__(self, inputs)
    169     """
    170     inputs = jnp.asarray(inputs, self.dtype)
--> 171     kernel = self.param('kernel',
    172                         self.kernel_init,
    173                         (inputs.shape[-1], self.features))

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/linen/module.py in param(self, name, init_fn, *init_args)
    783     if self._name_taken(name):
    784       raise errors.NameInUseError('param', name, self.__class__.__name__)
--> 785     v = self.scope.param(name, init_fn, *init_args)
    786     self._state.children[name] = 'params'
    787     return v

/nix/store/sn3d710d8mc362xxffhk1whjvbr10639-python3.9-flax-0.3.4/lib/python3.9/site-packages/flax/core/scope.py in param(self, name, init_fn, *init_args)
    617         # for inference to a half float type for example.
    618         if jnp.shape(val) != jnp.shape(abs_val):
--> 619           raise errors.ScopeParamShapeError(name, self.path_text, 
    620               jnp.shape(val), jnp.shape(abs_val))
    621     else:

ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "kernel" in "/Dense_0": (7, 12), (7, 13). (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamShapeError)

I know that I could zero out weights to achieve something similar, but I'm interested in benchmarking the actual wallclock difference between the large and small networks.

Relevant code appears to be here:

flax/flax/core/scope.py

Lines 607 to 624 in bbbbdcc

	if self.has_variable('params', name):
	abs_rng = jax.ShapeDtypeStruct((2,), jnp.uint32)
	value = self.get_variable('params', name)
	# Validate that the shape of the init_fn output is the same as the shape
	# of the existing parameter. This is to make sure that the hparams set up
	# in a Flax Module match the shapes coming in during apply, and if not,
	# catch it with an error message.
	# NOTE: We could consider moving this to `self.`
	abs_value = jax.eval_shape(lambda rng: init_fn(rng, *init_args), abs_rng)
	abs_value_flat = jax.tree_leaves(abs_value)
	value_flat = jax.tree_leaves(value)
	for val, abs_val in zip(value_flat, abs_value_flat):
	# NOTE: We could check dtype consistency here as well but it's
	# usefuleness is less obvious. We might intentionally change the dtype
	# for inference to a half float type for example.
	if jnp.shape(val) != jnp.shape(abs_val):
	raise errors.ScopeParamShapeError(name, self.path_text,
	jnp.shape(val), jnp.shape(abs_val))

.

How can I re-use the same network structure but adjust the shapes of intermediate arrays with flax?

[jheek]

You can change the shapes of the params that you pass to the model. However, Flax has a safety check to make sure that hyper params like the num_features passed to dense are consistent with the shapes of the params by passing the correct number of features to Dense:

class _net(nn.Module):
  @nn.compact
  def __call__(self, x, hidden_features):
    # A straightforward one hidden layer network...
    x = nn.Dense(hidden_features)(x)
    x = nn.relu(x)
    x = nn.Dense(5)(x)
    x = nn.log_softmax(x)
    return x
    
net.apply(init_params, jnp.zeros((1, 7)), hidden_features=13)

net.apply(sub_params, jnp.zeros((1, 7)), hidden_features=12)