Repository: google/jax

Issue:
vmap raises ambiguous exception when inconsistent sizes for array axes to be mapped occurs with user defined objects as static parameters
### Description

When trying to use vmap with functions, that accept other functions, an ambiguous error is raised if there is are inconsistent sizes for array axes to be mapped. A minimal example is shown below:

```
import jax.numpy as jnp
import jax


def functionThatTakesAFunction(array1,array2,function):
    """
    Takes in two arrays, applies function 1 to one of them and returns
    """
    array3 = function(array1)
    return jnp.vstack((array3,array2)),jnp.hstack((array3,array2))

def multiplyBy2(array):
    return 2 * array

batchMultiplyAndStack = jax.vmap(functionThatTakesAFunction,in_axes=[0,0,None])

test = jnp.array([[1,2],
                  [1,2]])
test2 = jnp.array([[1,2],
                   [1,2],
                   [1,2]])

testStacked = jnp.array((test,test,test))
test2Stacked = jnp.array((test2,test2,test2))

print(batchMultiplyAndStack(test,test2,multiplyBy2))
```

Running the above code produces the following stack trace:

```
Traceback (most recent call last):
  File "/Users/user/Desktop/jax/jax/_src/api_util.py", line 562, in shaped_abstractify
    return _shaped_abstractify_handlers[type(x)](x)
KeyError: <class 'function'>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/user/Desktop/test.py", line 41, in <module>
    print(batchMultiplyAndStack(test,test2,multiplyBy2))
  File "/Users/user/Desktop/jax/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1239, in vmap_f
    _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, "vmap"))
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1298, in _mapped_axis_size
    key_paths = [f'argument {name}{keystr(p)} '
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1299, in <listcomp>
    f'of type {_get_argument_type(x)}'
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1280, in _get_argument_type
    return shaped_abstractify(x).str_short()
  File "/Users/user/Desktop/jax/jax/_src/api_util.py", line 564, in shaped_abstractify
    return _shaped_abstractify_slow(x)
  File "/Users/user/Desktop/jax/jax/_src/api_util.py", line 553, in _shaped_abstractify_slow
    raise TypeError(
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Cannot interpret value of type <class 'function'> as an abstract array; it does not have a dtype attribute

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/user/Desktop/test.py", line 41, in <module>
    print(batchMultiplyAndStack(test,test2,multiplyBy2))
TypeError: Cannot interpret value of type <class 'function'> as an abstract array; it does not have a dtype attribute
```

I've tracked down the issue to how `jax._src.api._mapped_axis_size` creates the exception to inform users about inconsistent sizes for array axes to be mapped. When trying to get the data type of a function object, a TypeError is raised and not caught, resulting in an ambiguous exception as it is unclear why the exception is raised (since vmap normally works just fine with functions passed as static arguments).

I've made a pull request that fixes the issue (#16593), but I'm not sure if my fix (making a wrapper around a problematic call to  `shaped_abstractify(x).str_short()` is compliant with jax's coding standards. With this fix we get the expected stack trace for this error:

```
Traceback (most recent call last):
  File "/Users/user/Desktop/test.py", line 41, in <module>
    print(batchMultiplyAndStack(test,test2,multiplyBy2))
  File "/Users/user/Desktop/jax/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1239, in vmap_f
    _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, "vmap"))
  File "/Users/user/Desktop/jax/jax/_src/api.py", line 1322, in _mapped_axis_size
    raise ValueError(''.join(msg)[:-2])  # remove last semicolon and newline
jax._src.traceback_util.UnfilteredStackTrace: ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * one axis had size 2: axis 0 of argument array1 of type int32[2,2];
  * one axis had size 3: axis 0 of argument array2 of type int32[3,2]

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/user/Desktop/test.py", line 41, in <module>
    print(batchMultiplyAndStack(test,test2,multiplyBy2))
ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * one axis had size 2: axis 0 of argument array1 of type int32[2,2];
  * one axis had size 3: axis 0 of argument array2 of type int32[3,2]
```

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

jax v0.4.14 jaxlib v0.4.13

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

CPU

### Additional system info

Python 3.10.10, MacOS

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! Here's a shorter repro:
```python
import jax
import jax.numpy as jnp

def f(x1, x2, g):
  return g(x1, x2)

jax.vmap(f, (0, 0, None))(jnp.ones(2), jnp.ones(3), jnp.add)
```
It seems the main issue is that when you pass invalid arguments to `vmap`, it can fail in constructing the error message if there are non-array inputs.
That's a much more succinct way to put it! That's exactly the issue.

The root cause seems to be an uncaught exception when trying to determine the data type of the non-array inputs. The pull request I made tries to fix it. I just cleaned it up, but in squashing commits to one commit, there seems to be a noreply@google.com author (<no****ly​@google.com>) who's being blocked by the CLA check. Let me know if I need to do anything to fix that or if you'd rather fix the issue some other way.

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1275,6 +1275,11 @@ def _get_axis_size(name: str, shape: tuple[core.AxisSize, ...], axis: int
     msg = f"{name} must have at least one non-None value in in_axes"
     raise ValueError(msg)
 
+  def _get_argument_type(x):
+    try:
+      return shaped_abstractify(x).str_short()
+    except TypeError: #Catch all for user specified objects that can't be interpreted as a data type
+      return "unknown"
   msg = [f"{name} got inconsistent sizes for array axes to be mapped:\n"]
   args, kwargs = tree_unflatten(tree, vals)
   try:
@@ -1283,15 +1288,15 @@ def _get_axis_size(name: str, shape: tuple[core.AxisSize, ...], axis: int
     ba = None
   if ba is None:
     args_paths = [f'args{keystr(p)} '
-                  f'of type {shaped_abstractify(x).str_short()}'
+                  f'of type {_get_argument_type(x)}'
                   for p, x in generate_key_paths(args)]
     kwargs_paths = [f'kwargs{keystr(p)} '
-                    f'of type {shaped_abstractify(x).str_short()}'
+                    f'of type {_get_argument_type(x)}'
                     for p, x in generate_key_paths(kwargs)]
     key_paths = [*args_paths, *kwargs_paths]
   else:
     key_paths = [f'argument {name}{keystr(p)} '
-                 f'of type {shaped_abstractify(x).str_short()}'
+                 f'of type {_get_argument_type(x)}'
                  for name, arg in ba.arguments.items()
                  for p, x in generate_key_paths(arg)]
   all_sizes = [_get_axis_size(name, np.shape(x), d) if d is not None else None


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1699,6 +1699,15 @@ def test_device_put_sharding_mismatched_tree_different_leaf_count(self):
     ):
       jax.device_put((x, y, z), device=(s1, s2))
 
+  def test_vmap_inconsistent_sizes_constructs_proper_error_message(self):
+    def f(x1, x2, g):
+      return g(x1, x2)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "vmap got inconsistent sizes for array axes to be mapped:"
+    ):
+      jax.vmap(f, (0, 0, None))(jnp.ones(2), jnp.ones(3), jnp.add)
 
   def test_device_get_scalar(self):
     x = np.arange(12.).reshape((3, 4)).astype("float32")