Rollforward of Add a fastpath to pmap_lib for sharding np.ndarray dir…

…ectly in c++.

py::array::ensure(arg) was not a strict enough check and scalars were matching.
PiperOrigin-RevId: 513322424

tests/pmap_test.py

@@ -2115,6 +2115,12 @@ def test_cache_uses_jax_key(self):
     pmaped_f(inputs)
     self.assertEqual(pmaped_f._cache_size, 1)
 
+  def test_constants_fallback(self):
+    fn = pmap(lambda x, y: x + y, in_axes=(0, None))
+
+    for _ in range(2):
+      fn(np.zeros((jax.device_count(), 5), dtype=np.float32), 2.0)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class VmapOfPmapTest(jtu.JaxTestCase):