fonction with return type according to argument type but with intermediate change

[jecampagne]

Hello,
Finding a title is always a tricky business, as it has to be both concise and relevant enough, so I apologize for my poor.

The problem is to find a way to accept a scalar/vector as argument, turn it as a vector (at least 1d) to get an algorithm working, and finally return either a scalar or a vector according to the argument type. The function should also be jacrev compatible. Now a snippet that illustrate a possible solution (nb. the "algo" here would of course be thought to be complicated in real life, as a possible need to convert the scalar to vector think with a vmap):

def f(x):
  nd = jnp.ndim(x)
  x = jnp.atleast_1d(x)
  res = x**3    # here a very simple algo 
  if nd==0: 
    return res[0]
  else:
    return res

One can figure out that

xin = -12.; print(xin, " res=",f(xin))  # gives -12.0  res= -1728.0
xin = jnp.array([-12.]); print(xin, " res=",f(xin))  # gives [-12.]  res= [-1728.]
xin = -12.; print(jax.jacrev(f)(xin), 3*xin**2)  # gives 432.0 432.0
xin = jnp.array([-12.]); print(jax.jacrev(f)(xin), 3*xin**2)  # gives [[432.]] [432.]

So far so good, but my solution is certainly not the most elegant and robust. If you find better, I'd be grateful.

[jakevdp]

In internal JAX code, I often use a pattern that looks something like this:

def f(x):
  out_shape = jnp.shape(x)
  x = jnp.atleast_1d(x)
  res = x**3
  return res.reshape(out_shape)

[jecampagne]

I was not expecting more ! Thanks a lot.