fix test

stanford-crfm · Nov 24, 2024 · aa70b88 · aa70b88
1 parent c8553b4
commit aa70b88
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tests/test_histogram.py b/tests/test_histogram.py
@@ -11,19 +11,20 @@
 
 
 def test_sharded_histogram_simple():
-    mesh = Mesh((jax.devices()), (ResourceAxis.DATA))
+    mesh = Mesh((jax.devices()), (ResourceAxis.DATA,))
 
     Batch = hax.Axis("Batch", 64)
     Feature = hax.Axis("Feature", 128)
 
     with mesh, hax.axis_mapping({"batch": ResourceAxis.DATA}):
         a = hax.random.normal(PRNGKey(0), (Batch, Feature))
         a = hax.shard(a)
-        hist = levanter.tracker.histogram.sharded_histogram(a, bins=10)
+        hist, bins = levanter.tracker.histogram.sharded_histogram(a, bins=10)
 
-    hist_normal = jax.numpy.histogram(a.array, bins=10)[0]
+    hist_normal, bins_normal = jax.numpy.histogram(a.array, bins=10)
 
     assert jax.numpy.allclose(hist, hist_normal)
+    assert jax.numpy.allclose(bins, bins_normal)
 
 
 @skip_if_not_enough_devices(2)