Skip to content

Commit

Permalink
[JAX] Update JAX users in preparation for a change that makes iterati…
Browse files Browse the repository at this point in the history
…on over a JAX array return JAX arrays, instead of NumPy arrays.

See #8043 for context as to why we are making this change.

The upshot for most users is that the values returned by iteration over a JAX array are now themselves JAX arrays, with the semantics of JAX arrays, which sometimes differ from the semantics of NumPy scalars and arrays.

PiperOrigin-RevId: 405780628
  • Loading branch information
sharadmv authored and jax authors committed Oct 27, 2021
1 parent b8296bb commit 259f5a0
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions tests/ann_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def test_approx_max_k(self, qy_shape, db_shape, dtype, k, recall):
_, gt_args = lax.top_k(scores, k)
_, ann_args = ann.approx_max_k(scores, k, recall_target=recall)
self.assertEqual(k, len(ann_args[0]))
gt_args_sets = [set(x) for x in gt_args]
gt_args_sets = [set(np.asarray(x)) for x in gt_args]
hits = sum(
len(list(x
for x in ann_args_per_q
Expand Down Expand Up @@ -77,7 +77,7 @@ def test_approx_min_k(self, qy_shape, db_shape, dtype, k, recall):
_, gt_args = lax.top_k(-scores, k)
_, ann_args = ann.approx_min_k(scores, k, recall_target=recall)
self.assertEqual(k, len(ann_args[0]))
gt_args_sets = [set(x) for x in gt_args]
gt_args_sets = [set(np.asarray(x)) for x in gt_args]
hits = sum(
len(list(x
for x in ann_args_per_q
Expand Down

0 comments on commit 259f5a0

Please sign in to comment.