Fix batch size calculation with vmap and static_argnums

[dime10]

Resolves a bug accessing argument shapes in the vmap function when those arguments don't have a shape. This might happen when using vmap together with static_argnums for example:

@qjit(static_argnums=0)
@vmap(in_axes=(None, 0))
def test(n, array):
    for _ in range(n):
        array = array * 2
    return array

test(10, np.random.rand(3,2))

Also improves error message for out of bounds axis specifier.

[sc-74291]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.87%. Comparing base (c13cd9b) to head (8b23dda).
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1150      +/-   ##
==========================================
- Coverage   97.88%   97.87%   -0.02%     
==========================================
  Files          76       76              
  Lines       10863    10850      -13     
  Branches     1283     1283              
==========================================
- Hits        10633    10619      -14     
  Misses        179      179              
- Partials       51       52       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.