Squashed commit of the following:

commit 58f05dd
Merge: a8ba0c7 f4891cf
Author: Kaya Unalmis <kayaunalmis@proton.me>
Date:   Tue Aug 19 19:52:38 2025 -0500

    Merge branch 'master' into ku/shard_0

commit a8ba0c7
Author: unalmis <kayaunalmis@proton.me>
Date:   Tue Aug 19 19:48:28 2025 -0500

    Remove test

commit f8324cb
Author: unalmis <kayaunalmis@proton.me>
Date:   Tue Aug 19 19:44:30 2025 -0500

    Pulling progress from ku/shard on #1773

Author: unalmis

desc/batching.py

@@ -283,8 +283,7 @@ def _evaluate_in_chunks(
             for i, a in enumerate(args)
         ]
     )
-    # Note that ``num_batches`` in ``_batch_and_remainder``
-    # is ``n_elements//chunk_size`` which is always positive.
+    # Note that num_batches in _batch_and_remainder is always positive.
     scan_y = _scanmap(vmapped_fun, argnums, reduction, chunk_reduction)(
         *scan_x, **kwargs
     )