Repository: google/flax

Issue:
apply_gradient with no parameters gives ValueError
This issue is admittedly a corner case, but one we've run into. If we consider the following `flax.nn.Module`:

```python
class Identity(flax.nn.Module):
    def apply(self, x):
        return x
```

We won't be able to call `apply_gradient` since the output from [this line](https://github.com/google/flax/blob/master/flax/optim/base.py#L134) will be an empty list.

This should probably (?) be addressed since it's exceptional behavior that may surprise, but could see arguments for different ways of resolving. One simple answer is to just no-op, but there might be some higher-level concerns I'm not thinking about which say we don't even want parameterless modules (in which case, raise on construction).

Anyway, we've resolved for now by just adding a dummy parameter. Here's the full minimum example and the resulting value error:

```python
import flax
import jax
import jax.numpy as jnp

class Identity(flax.nn.Module):
    def apply(self, x):
        return x

model_def = Identity.partial()
_, params = model_def.init_by_shape(jax.random.PRNGKey(0), [(1,)])
model = flax.nn.Model(model_def, params)

def loss_fn(model, x, y):
    y_hat = model(x)
    return jnp.square(y - y_hat).mean(), y_hat

optim_def = flax.optim.Adam(learning_rate=1.0)
optimizer = optim_def.create(model)

(loss, y_hat), grad = jax.value_and_grad(loss_fn, has_aux=True)(optimizer.target, 1.0, 2.0)
optimizer.apply_gradient(grad)
```

```python
~/src/flax/flax/optim/base.py in apply_gradient(self, hyper_params, params, state, grads)
    135            for param, state, grad in zip(params_flat, states_flat, grads_flat)]
    136 
--> 137     new_params_flat, new_states_flat = list(zip(*out))
    138     new_params = jax.tree_unflatten(treedef, new_params_flat)
    139     new_param_states = jax.tree_unflatten(treedef, new_states_flat)

ValueError: not enough values to unpack (expected 2, got 0)
```


Discussion:

Resolving patch:
diff --git a/flax/optim/base.py b/flax/optim/base.py
--- a/flax/optim/base.py
+++ b/flax/optim/base.py
@@ -134,7 +134,7 @@ def apply_gradient(self, hyper_params, params, state, grads):
     out = [self.apply_param_gradient(step, hyper_params, param, state, grad)
            for param, state, grad in zip(params_flat, states_flat, grads_flat)]
 
-    new_params_flat, new_states_flat = list(zip(*out))
+    new_params_flat, new_states_flat = list(zip(*out)) if out else ((), ())
     new_params = jax.tree_unflatten(treedef, new_params_flat)
     new_param_states = jax.tree_unflatten(treedef, new_states_flat)
     new_state = OptimizerState(step + 1, new_param_states)


Test patch:
diff --git a/tests/optim_test.py b/tests/optim_test.py
--- a/tests/optim_test.py
+++ b/tests/optim_test.py
@@ -79,6 +79,14 @@ def test_optimizer_with_focus(self):
     self.assertEqual(new_optimizer.state, expected_state)
     self.assertEqual(new_optimizer.target, expected_params)
 
+  def test_empty_optimizer(self):
+    params = {}
+    optimizer_def = optim.Momentum(learning_rate=0.1)
+    optimizer = optimizer_def.create(params)
+    new_optimizer = optimizer.apply_gradient({})
+    expected_state = optim.OptimizerState(1, {})
+    self.assertEqual(new_optimizer.state, expected_state)
+
 
 class ModelParamTraversalTest(absltest.TestCase):