Require array-like inputs to sparse_plus

We should not silently convert non-array inputs to arrays, because this can lead to silent performance degredation. This brings the sparse_plus API in line with other APIs in this module.

PiperOrigin-RevId: 617190413

Author: Jake VanderPlas

jax/_src/nn/functions.py

@@ -132,6 +132,7 @@ def sparse_plus(x: ArrayLike) -> Array:
   Args:
     x: input (float)
   """
+  numpy_util.check_arraylike("sparse_plus", x)
   x = jnp.asarray(x)
   return jnp.where(x <= -1.0, 0.0, jnp.where(x >= 1.0, x, (x + 1.0)**2/4))