Repository: google/jax

Issue:
Vectorizing over None argument returns NaNs
### Description

Hi,

I would like to do the following thing:

```python
import jax.numpy as jnp

def f(x, y):
    if y is None:
        return x
    return x * y

f = jnp.vectorize(f, signature="(),()->()")
```
however, right now, if you pass it a `None` as the second argument, the returned value is `NaN` (or an array thereof).

```python
f(jnp.ones(6), None)  # Returns Array([nan, nan, nan, nan, nan, nan], dtype=float32)
f(1., None)  # Returns Array(nan, dtype=float32)
```

I don't think this is a valid behaviour and the call should probably either fail explicitly (like `jax.vmap`) or return the "expected" value.

WDYT?

### What jax/jaxlib version are you using?

jax v0.4.13

### Which accelerator(s) are you using?

CPU

### Additional system info

Google Colab

### NVIDIA GPU info

_No response_


Discussion:
`y is None` will always be false for traced values, because `None` is a Python object and traced values will be `jax.Array` objects. When you vectorize a function, ~it works via `vmap`, which converts all mapped inputs to arrays.~ it converts all inputs to arrays before passing them to `vmap`.

If you want to use statements like `y is None` in your code, you need to make sure that `y` never participates in a JAX transformation. In your case, it might look like this:
```python
f_wrapped = lambda x, y: jnp.vectorize(partial(f, y=y), signature="()->()")(x)

f_wrapped(jnp.ones(6), None)
# Array([1., 1., 1., 1., 1., 1.], dtype=float32)
```

@jakevdp I'm surprised to hear this, since the following code works:

```python
@jax.jit
@partial(jax.vmap, in_axes=(0, 0))
def f(x, y):
    if y is None:
        return x

    return x + y

print(f(jnp.arange(3), None))          # [0, 1, 2]                                     
print(f(jnp.arange(3), jnp.arange(3))) # [0, 2, 4] 
```

My mental model had been that `None` is an empty pytree, so when it needs to be traced it gets convert to an empty pytree of tracers. Is this officially unsupported behavior that might change at some point?
Oh, thanks for pointing that out – I think the `None` conversion is not a feature of `jnp.vmap`, but rather a feature of `jnp.vectorize`, which does not accept general pytree arguments as inputs. The reason for this difference is that the `jnp.vectorize` API is modeled after `np.vectorize` and NumPy does not have any notion of pytrees, so it treats all inputs to a vectorized function as arrays.
Ah okay, makes sense. I was worried I was about to need to change a _lot_ of code :sweat_smile: 
For what it's worth, your code doesn't work with `np.vectorize` either, becuase it also converts `y` to an array, so that `y is None` will return False:
```python
In [1]: import numpy as np
   ...: 
   ...: def f(x, y):
   ...:     if y is None:
   ...:         return x
   ...:     return x * y
   ...: 
In [2]: f = np.vectorize(f, signature="(),()->()")

In [3]: f(None, 1)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[3], line 1
----> 1 f(None, 1)

File ~/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.9/site-packages/numpy/lib/function_base.py:2372, in vectorize.__call__(self, *args, **kwargs)
   2369     self._init_stage_2(*args, **kwargs)
   2370     return self
-> 2372 return self._call_as_normal(*args, **kwargs)

File ~/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.9/site-packages/numpy/lib/function_base.py:2365, in vectorize._call_as_normal(self, *args, **kwargs)
   2362     vargs = [args[_i] for _i in inds]
   2363     vargs.extend([kwargs[_n] for _n in names])
-> 2365 return self._vectorize_call(func=func, args=vargs)

File ~/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.9/site-packages/numpy/lib/function_base.py:2446, in vectorize._vectorize_call(self, func, args)
   2444 """Vectorized call to `func` over positional `args`."""
   2445 if self.signature is not None:
-> 2446     res = self._vectorize_call_with_signature(func, args)
   2447 elif not args:
   2448     res = func()

File ~/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.9/site-packages/numpy/lib/function_base.py:2486, in vectorize._vectorize_call_with_signature(self, func, args)
   2483 nout = len(output_core_dims)
   2485 for index in np.ndindex(*broadcast_shape):
-> 2486     results = func(*(arg[index] for arg in args))
   2488     n_results = len(results) if isinstance(results, tuple) else 1
   2490     if nout != n_results:

Cell In[29], line 6, in f(x, y)
      4 if y is None:
      5     return x
----> 6 return x * y

TypeError: unsupported operand type(s) for *: 'NoneType' and 'int'
```
So I think `jax.numpy.vectorize` is arguably working as expected here.
You got the order of the arguments wrong, it would be f(1, None), which works fine with numpy.
Sorry I lost track of this, I think #18441 should fix the issue. Here are the results with that change:
```python
In [1]: import jax.numpy as jnp
   ...: 
   ...: def f(x, y):
   ...:     if y is None:
   ...:         return x
   ...:     return x * y
   ...: 
   ...: f = jnp.vectorize(f, signature="(),()->()")

In [2]: f(jnp.ones(6), None)
Out[2]: Array([1., 1., 1., 1., 1., 1.], dtype=float32)

In [3]: f(1., None)
Out[3]: Array(1., dtype=float32, weak_type=True)
```

Resolving patch:
diff --git a/jax/_src/numpy/vectorize.py b/jax/_src/numpy/vectorize.py
--- a/jax/_src/numpy/vectorize.py
+++ b/jax/_src/numpy/vectorize.py
@@ -263,7 +263,6 @@ def wrapped(*args):
     error_context = ("on vectorized function with excluded={!r} and "
                      "signature={!r}".format(excluded, signature))
     excluded_func, args = _apply_excluded(pyfunc, excluded, args)
-    args = tuple(map(jnp.asarray, args))
 
     if signature is not None:
       input_core_dims, output_core_dims = _parse_gufunc_signature(signature)
@@ -271,6 +270,15 @@ def wrapped(*args):
       input_core_dims = [()] * len(args)
       output_core_dims = None
 
+    none_args = {i for i, arg in enumerate(args) if arg is None}
+    if any(none_args):
+      if any(input_core_dims[i] != () for i in none_args):
+        raise ValueError(f"Cannot pass None at locations {none_args} with {signature=}")
+      excluded_func, args = _apply_excluded(excluded_func, none_args, args)
+      input_core_dims = [dim for i, dim in enumerate(input_core_dims) if i not in none_args]
+
+    args = tuple(map(jnp.asarray, args))
+
     broadcast_shape, dim_sizes = _parse_input_dimensions(
         args, input_core_dims, error_context)
 


Test patch:
diff --git a/tests/lax_numpy_vectorize_test.py b/tests/lax_numpy_vectorize_test.py
--- a/tests/lax_numpy_vectorize_test.py
+++ b/tests/lax_numpy_vectorize_test.py
@@ -225,6 +225,22 @@ def test_expand_dims_multiple_outputs_no_signature(self):
     self.assertAllClose(xx[1], x)
     self.assertIsInstance(xx, tuple)
 
+  def test_none_arg(self):
+    f = jnp.vectorize(lambda x, y: x if y is None else x + y)
+    x = jnp.arange(10)
+    self.assertAllClose(f(x, None), x)
+
+    y = jnp.arange(10, 20)
+    self.assertAllClose(f(x, y), x + y)
+
+  def test_none_arg_bad_signature(self):
+    f = jnp.vectorize(lambda x, y: x if y is None else x + y,
+                      signature='(k),(k)->(k)')
+    args = jnp.arange(10), None
+    msg = r"Cannot pass None at locations \{1\} with signature='\(k\),\(k\)->\(k\)'"
+    with self.assertRaisesRegex(ValueError, msg):
+      f(*args)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())