Add tests for mean and var with symbolic shape

test/test_symbolic_jit.py

@@ -177,5 +177,27 @@ def f(a): return (a+1).realize()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
     assert_jit_cache_len(jf, 1)
 
+  def test_mean(self):
+    def f(a): return a.mean().realize()
+    jf = TinyJit(f)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(3, i)
+      symbolic = jf(a.reshape(3, vi)).numpy()
+      expected = f(a).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 1)
+
+  def test_var(self):
+    def f(a): return a.var().realize()
+    jf = TinyJit(f)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(3, i)
+      symbolic = jf(a.reshape(3, vi)).numpy()
+      expected = f(a).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+    assert_jit_cache_len(jf, 2)
+
 if __name__ == '__main__':
   unittest.main()

test/test_symbolic_ops.py

@@ -141,5 +141,27 @@ def test_shrink(self):
       expected = a.shrink(((3,5),(i,i+2))).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
 
+  def test_mean(self):
+    def f(a): return a.mean(axis=1).realize()
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(3, i)
+      symbolic = f(a.reshape(3, vi)).numpy()
+      expected = f(a).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      var_vals = a.reshape(3, vi).mean().schedule_with_vars()[1]
+      self.assertEqual(var_vals[vi.unbind()[0]], i)
+
+  def test_var(self):
+    def f(a): return a.var(axis=1).realize()
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      a = Tensor.rand(3, i)
+      symbolic = f(a.reshape(3, vi)).numpy()
+      expected = f(a).numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+      var_vals = a.reshape(3, vi).var().schedule_with_vars()[1]
+      self.assertEqual(var_vals[vi.unbind()[0]], i)
+
 if __name__ == '__main__':
   unittest.main()