Repository: google/jax

Issue:
2nd-order autodiff of remat of `dot_general` with `float` input leads to error
This came up in the process of creating #16721

Here's a minimal repro:
```python
import jax

@jax.remat
def f(x):
  return jax.lax.dot_general(x, x, (([], []), ([], [])))

jax.hessian(f)(1.0)
```
```pytb
Traceback (most recent call last):
  File "/Users/vanderplas/github/google/jax/tmp.py", line 7, in <module>
    jax.hessian(f)(1.0)
  File "/Users/vanderplas/github/google/jax/tmp.py", line 5, in f
    return jax.lax.dot_general(x, x, (([], []), ([], [])))
jax._src.source_info_util.JaxStackTraceBeforeTransformation: AttributeError: 'float' object has no attribute 'ndim'
```

<details>
<summary>Show full traceback</summary>

```pytb
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/vanderplas/github/google/jax/tmp.py", line 7, in <module>
    jax.hessian(f)(1.0)
  File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/api.py", line 1257, in vmap_f
    out_flat = batching.batch(
               ^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/linear_util.py", line 188, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/api.py", line 1986, in _jvp
    out_primals, out_tangents = ad.jvp(flat_fun).call_wrapped(ps_flat, ts_flat)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/linear_util.py", line 188, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/api.py", line 946, in jacfun
    y, pullback = _vjp(f_partial, *dyn_args)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/api.py", line 2235, in _vjp
    out_primal, out_vjp = ad.vjp(
                          ^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 139, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 128, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/partial_eval.py", line 776, in trace_to_jaxpr_nounits
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
                                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/linear_util.py", line 188, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/ad_checkpoint.py", line 285, in fun_remat
    out_flat = remat_p.bind(
               ^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 379, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 382, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 315, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/ad_checkpoint.py", line 475, in remat_jvp
    outs = remat_p.bind(
           ^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 379, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 382, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/partial_eval.py", line 212, in process_primitive
    return custom_partial_eval_rules[primitive](self, *tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/ad_checkpoint.py", line 513, in remat_partial_eval
    out_consts = core.eval_jaxpr(jaxpr_known, (), *in_consts)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 447, in eval_jaxpr
    ans = eqn.primitive.bind(*subfuns, *map(read, eqn.invars), **bind_params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 379, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 382, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 315, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 531, in standard_jvp
    tangents_out = [rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 531, in <listcomp>
    tangents_out = [rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/ad.py", line 557, in <lambda>
    lhs_jvp = lambda g, x, y, **kwargs: prim.bind(g, y, **kwargs)
                                        ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 379, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/core.py", line 382, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/interpreters/batching.py", line 430, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vanderplas/github/google/jax/jax/_src/lax/lax.py", line 2610, in _dot_general_batch_rule
    (lhs.ndim, rhs.ndim), (left_stack_dim, right_stack_dim),
               ^^^^^^^^
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: 'float' object has no attribute 'ndim'

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/vanderplas/github/google/jax/tmp.py", line 7, in <module>
    jax.hessian(f)(1.0)
  File "/Users/vanderplas/github/google/jax/jax/_src/api.py", line 859, in jacfun
    y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'float' object has no attribute 'ndim'
```

</details>


Discussion:
I think `_dot_general_batch_rule` should accept Python numeric types for `lhs` and `rhs`, i.e. it shouldn't expect the `ndim` attribute to exist on them and should instead use `np.ndim` to handle the polymorphism.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2608,7 +2608,7 @@ def _dot_general_batch_rule(batched_args, batch_dims, *, dimension_numbers,
   left_stack_dim = lbd.stacked_axis if type(lbd) is RaggedAxis else lbd
   right_stack_dim = rbd.stacked_axis if type(rbd) is RaggedAxis else rbd
   new_dimension_numbers, result_stack_dim = _dot_general_batch_dim_nums(
-      (lhs.ndim, rhs.ndim), (left_stack_dim, right_stack_dim),
+      (np.ndim(lhs), np.ndim(rhs)), (left_stack_dim, right_stack_dim),
       dimension_numbers)
   # TODO Should probably check that any ragged dimensions have corresponding
   # sizes, because otherwise the dot product is technically undefined.
@@ -2619,12 +2619,12 @@ def _dot_general_batch_rule(batched_args, batch_dims, *, dimension_numbers,
     lhs = batching.mask_ragged_axes(lhs, _get_sum_identity, lbd)
     lhs_shape = batching.bdim_as_shape(lbd, lhs.shape)
   else:
-    lhs_shape = lhs.shape
+    lhs_shape = np.shape(lhs)
   if type(rbd) is RaggedAxis:
     rhs = batching.mask_ragged_axes(rhs, _get_sum_identity, rbd)
     rhs_shape = batching.bdim_as_shape(rbd, rhs.shape)
   else:
-    rhs_shape = rhs.shape
+    rhs_shape = np.shape(rhs)
   batched_out = dot_general(lhs, rhs, new_dimension_numbers,
                             precision=precision,
                             preferred_element_type=preferred_element_type)


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2664,6 +2664,14 @@ def testDynamicSliceU8Index(self):
     np.testing.assert_equal(
         np.array(lax.dynamic_slice(x, np.uint8([128]), (1,))), [128])
 
+  def test_dot_general_batching_python_builtin_arg(self):
+    # https://github.com/google/jax/issues/16805
+    @jax.remat
+    def f(x):
+      return jax.lax.dot_general(x, x, (([], []), ([], [])))
+
+    jax.hessian(f)(1.0)  # don't crash
+
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):