Repository: google/jax
Issue: sm3 optimizer doesnt work with scalar data

This is almost tutorial code, it doesnt work with sm3 because of scalar data

```python
import jax
import jax.numpy as jnp
import numpy as np

from typing import NamedTuple
import functools

class Params(NamedTuple):
  weight: jnp.ndarray
  bias: jnp.ndarray


def init(rng) -> Params:
  """Returns the initial model params."""
  weights_key, bias_key = jax.random.split(rng)
  weight = jax.random.normal(weights_key, ())
  bias = jax.random.normal(bias_key, ())
  return Params(weight, bias)


def loss_fn(params: Params, xs: jnp.ndarray, ys: jnp.ndarray) -> jnp.ndarray:
  """Computes the least squares error of the model's predictions on x against y."""
  pred = params.weight * xs + params.bias
  return jnp.mean((pred - ys) ** 2)

LEARNING_RATE = 0.005

true_w, true_b = 2, -1
xs = np.random.normal(size=(128, 1))
noise = 0.5 * np.random.normal(size=(128, 1))
ys = xs * true_w + true_b + noise

# Initialise parameters and replicate across devices.
params = init(jax.random.PRNGKey(123))

from jax.experimental.optimizers import sm3, adam

opt_init, opt_update, get_params = sm3(step_size=LEARNING_RATE)  
opt_state = opt_init(params)

def update(params: Params, opt_state, xs: jnp.ndarray, ys: jnp.ndarray, i) -> Params:
    """Performs one SGD update step on params using the given data."""
    # Compute the gradients on the given minibatch (individually on each device).
    loss, grads = jax.value_and_grad(loss_fn)(params, xs, ys)
    # Each device performs its own update, but since we start with the same params
    # and synchronise gradients, the params stay in sync.

    opt_state = opt_update(i, grads, opt_state)
    params = get_params(opt_state)

    return params, opt_state, loss

def type_after_update(name, obj):
  print(f"after first `update()`, `{name}` is a", type(obj))

# Actual training loop.
for i in range(1000):

  # This is where the params and data gets communicated to devices:
  params, opt_state, loss = update(params, opt_state, xs, ys, i)

  # The returned `replicated_params` and `loss` are now both ShardedDeviceArrays,
  # indicating that they're on the devices.
  # `x_split`, of course, remains a NumPy array on the host.
  if i == 0:
    type_after_update('replicated_params.weight', params.weight)
    type_after_update('loss', loss)
    type_after_update('x_split', xs)

  if i % 100 == 0:
    # Note that loss is actually an array of shape [num_devices], with identical
    # entries, because each device returns its copy of the loss.
    # So, we take the first element to print it.
    print(f"Step {i:3d}, loss: {loss:.3f}")

```

The issue is in this line of code in sm3 optimizer 
```
vs = [jnp.zeros(sz, dtype=x0.dtype) for sz in x0.shape]
```

If x0 is scalar then it will generate an empty array. Its possible to fix that with something like this 
```
  def s2v(s):
    return s if s.ndim > 0 else jnp.array([s])

  def init(x0):
    vs = [jnp.zeros(sz, dtype=x0.dtype) for sz in s2v(x0).shape]
    return x0, jnp.zeros_like(x0), vs
```

but it will lead to the output with 1-element arrays instead of scalars. 



Changed files:
tests/optimizers_test.py, jax/experimental/optimizers.py

diff --git a/jax/experimental/optimizers.py b/jax/experimental/optimizers.py
index aaaaef8f4..efea53a97 100644
--- a/jax/experimental/optimizers.py
+++ b/jax/experimental/optimizers.py
@@ -478,22 +478,24 @@ def sm3(step_size, momentum=0.9):
     return x[tuple(idx)]
 
   def init(x0):
+    x_shape = x0.shape
+    x0 = jnp.atleast_1d(x0)
     vs = [jnp.zeros(sz, dtype=x0.dtype) for sz in x0.shape]
-    return x0, jnp.zeros_like(x0), vs
+    return x0, jnp.zeros_like(x0), vs, x_shape
 
   def update(i, g, state):
-    x, m, vs = state
+    x, m, vs, x_shape = state
     vs = [broadcast_into(g.ndim, v, i) for i, v in enumerate(vs)]
     accum = functools.reduce(jnp.minimum, vs) + jnp.square(g)
     accum_inv_sqrt = jnp.where(accum > 0, 1. / jnp.sqrt(accum), 0)
     m = (1. - momentum) * (g * accum_inv_sqrt) + momentum * m
     x = x - step_size(i) * m
     vs = [accum.max(splice(range(x.ndim), j, [])) for j in range(x.ndim)]
-    return x, m, vs
+    return x, m, vs, x_shape
 
   def get_params(state):
-    x, _, _ = state
-    return x
+    x, _, _, x_shape = state
+    return x.reshape(x_shape)
 
   return init, update, get_params
 
diff --git a/tests/optimizers_test.py b/tests/optimizers_test.py
index ce7ee4dc0..9819e9165 100644
--- a/tests/optimizers_test.py
+++ b/tests/optimizers_test.py
@@ -138,7 +138,14 @@ class OptimizerTests(jtu.JaxTestCase):
     x0 = (jnp.ones(2), jnp.ones((2, 2)))
     self._CheckOptimizer(optimizers.adagrad, loss, x0, num_iters, step_size)
 
-  def testSM3(self):
+  def testSM3Scalar(self):
+    def loss(x): return x**2
+    x0 = jnp.array(1.)
+    num_iters = 100
+    step_size = 0.1
+    self._CheckOptimizer(optimizers.sm3, loss, x0, num_iters, step_size)
+
+  def testSM3Vector(self):
     def loss(xs):
       x1, x2 = xs
       return jnp.sum(x1 ** 2) + jnp.sum(x2 ** 2)