Autoconvert from float to jax.numpy

[patel-zeel]

I have recently started working with jaxopt and I found that jaxopt does not auto-convert float to jax.numpy array. For example, If I run the following code:

Example 1

import jax
import jax.numpy as jnp
import jaxopt
import optax

def loss_function(params):
    return (params-1)**2

params = 0.0
jaxopt.ScipyMinimize(fun=loss_function).run(params)

It fails because params is just a float, but if I change last two lines with the following:

params = jnp.array(0.0)
jaxopt.ScipyMinimize(fun=loss_function).run(params)

The above code works fine. On the other side, optax does not fail in this case (maybe because jax.grad auto converts float to jax.numpy). Consider the following example for the same:

Example 2

params = 0.0
tx = optax.adam(learning_rate=0.1)
state = tx.init(params)
grads = jax.grad(loss_function)(params)
updates, state = tx.update(grads, state)
params = optax.apply_updates(params, updates)
print(params)

Output

0.09999933

So, I was wondering if jaxopt should auto-convert floats to jax.numpy or will that create some new issues?

Edit:
Error trace shows the following lines:

File ~/miniconda3/envs/jax_cpu/lib/python3.8/site-packages/jaxopt/_src/scipy_wrappers.py:192, in <listcomp>(.0)
    190 """Returns a PyTreeTopology encoding the PyTree structure of `x_jnp`."""
    191 leaves, treedef = tree_util.tree_flatten(x_jnp)
--> 192 shapes = [leaf.shape for leaf in leaves]
    193 dtypes = [leaf.dtype for leaf in leaves]
    194 return PyTreeTopology(treedef=treedef, shapes=shapes, dtypes=dtypes)

AttributeError: 'float' object has no attribute 'shape'

Will it suffice to just replace a few lines like this? (I tried it locally and it works for Example 1)

-shapes = [leaf.shape for leaf in leaves]
-dtypes = [leaf.dtype for leaf in leaves]

+shapes = [() if isinstance(leaf, float) else leaf.shape for leaf in leaves]
+default_dtype = jnp.float64 if jax.config.x64_enabled else jnp.float32
+dtypes = [default_dtype if isinstance(leaf, float) else leaf.dtype for leaf in leaves]