* fix test

debug/vmap_callback.py

@@ -0,0 +1,20 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+
+def add(x, y):
+    print(x.shape, y.shape)
+    return x + y
+
+@jax.jit
+@partial(jax.vmap, in_axes=(0, None))
+def cb(x, y):
+    return jax.pure_callback(add, jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype), x, y, vectorized=False)
+
+
+if __name__ == '__main__':
+    x = jnp.arange(4, dtype=jnp.float32)
+    y = jnp.asarray(1.)
+    print(cb(x, y))

dsa2000_cal/dsa2000_cal/visibility_model/source_models/celestial/tests/test_point_source_model.py

@@ -135,8 +135,8 @@ def test_gh55_point():
     pixsize = 0.5 * np.pi / 180 / 3600.  # 5 arcsec
     x0 = 0.
     y0 = 0.
-    l0 = y0
-    m0 = x0
+    l0 = x0
+    m0 = y0
     dl = pixsize
     dm = pixsize
     dirty = np.zeros((N, N))  # [Nl, Nm]
@@ -175,8 +175,8 @@ def pixel_to_lmn(xi, yi):
         dirty=dirty,
         pixsize_m=dm,
         pixsize_l=dl,
-        center_m=x0,
-        center_l=y0,
+        center_m=m0,
+        center_l=l0,
         epsilon=1e-4
     )
     print(vis)
@@ -214,7 +214,7 @@ def pixel_to_lmn(xi, yi):
     image = np.zeros((2, num_freqs, 2, 2))  # [source, chan, 2, 2]
     image[:, :, 0, 0] = 0.5
     image[:, :, 1, 1] = 0.5
-    gains = np.ones((2, num_ants, num_freqs, 2, 2), jnp.complex64)  # [[source,] time, ant, chan, 2, 2]
+    gains = np.zeros((2, num_ants, num_freqs, 2, 2), jnp.complex64)  # [[source,] time, ant, chan, 2, 2]
     gains[..., 0, 0] = 1.
     gains[..., 1, 1] = 1.
     lmn = np.stack([lmn1, lmn2], axis=0)  # [source, 3]