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_

Changed files:
tests/api_test.py, jax/_src/api.py

diff --git a/jax/_src/api.py b/jax/_src/api.py
index c9f3755ee..50d65dc14 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1275,6 +1275,11 @@ def _mapped_axis_size(fn, tree, vals, dims, name):
     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 _mapped_axis_size(fn, tree, vals, dims, name):
     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
diff --git a/tests/api_test.py b/tests/api_test.py
index 039abca31..a87671033 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1699,6 +1699,15 @@ class APITest(jtu.JaxTestCase):
     ):
       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")