fix(imperative/python): fix warp_perspective doc and arange dtype

GitOrigin-RevId: dc4e021ec0

imperative/python/megengine/functional/elemwise.py

@@ -487,8 +487,6 @@ def clip(x: Tensor, lower=None, upper=None) -> Tensor:
     ), "At least one of 'lower' or 'upper' must not be None"
     if lower is not None:
         if upper is not None:
-            # FIXME: following assertion won't work during trace if upper and lower are Tensors
-            # assert lower <= upper, "clip lower bound is bigger that upper bound"
             return minimum(maximum(x, lower), upper)
         else:
             return maximum(x, lower)

imperative/python/megengine/functional/tensor.py

@@ -1063,7 +1063,7 @@ def linspace(
 
     op = builtin.Linspace(comp_node=device)
     (result,) = apply(op, start, stop, num)
-    if np.dtype(dtype) == np.int32:
+    if np.dtype(dtype) != np.float32:
         return result.astype(dtype)
     return result
 
@@ -1112,7 +1112,7 @@ def arange(
     num = ceil((stop - start) / step)
     stop = start + step * (num - 1)
     result = linspace(start, stop, num, device=device)
-    if np.dtype(dtype) == np.int32:
+    if np.dtype(dtype) != np.float32:
         return result.astype(dtype)
     return result
 

imperative/python/megengine/functional/vision.py

@@ -398,8 +398,8 @@ def warp_perspective(
 
     .. math::
             \text{output}(n, c, h, w) = \text{input} \left( n, c,
-                \frac{M_{00}h + M_{01}w + M_{02}}{M_{20}h + M_{21}w + M_{22}},
-                \frac{M_{10}h + M_{11}w + M_{12}}{M_{20}h + M_{21}w + M_{22}}
+                \frac{M_{00}w + M_{01}h + M_{02}}{M_{20}w + M_{21}h + M_{22}},
+                \frac{M_{10}w + M_{11}h + M_{12}}{M_{20}w + M_{21}h + M_{22}}
                 \right)
 
     Optionally, we can set `mat_idx` to assign different transformations to the same image,

imperative/python/megengine/jit/tracing.py

@@ -1041,12 +1041,6 @@ class trace:
             raise RuntimeError("trace is not set with profiling=True")
         return json.loads(self._profiler.get())
 
-    def trace(self, *args, **kwargs):
-        raise NotImplementedError(
-            "trace is deemed unbeneficial with the new "
-            "tracing mechanism. You should alwasy use __call__."
-        )
-
 
 class CompiledTensorProxy:
     r"""Duck-typed RawTensor"""

imperative/python/test/run.sh

@@ -34,18 +34,18 @@ if [[ "$TEST_PLAT" =~ "local" ]]; then
     esac
 
     echo "test local env at: ${test_dirs}"
-    PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -v $test_dirs -m 'not isolated_distributed'
+    PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'not isolated_distributed'
     if [[ "$TEST_PLAT" =~ "cuda" ]]; then
         echo "test GPU pytest now"
-        PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -v $test_dirs -m 'isolated_distributed'
+        PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'isolated_distributed'
     fi
 else
     cd $(dirname "${BASH_SOURCE[0]}")/..
     test_dirs="megengine test"
     echo "test develop env"
-    PYTHONPATH="." PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -v $test_dirs -m 'not isolated_distributed'
+    PYTHONPATH="." PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'not isolated_distributed'
     if [[ "$TEST_PLAT" =~ "cuda" ]]; then
         echo "test GPU pytest now"
-        PYTHONPATH="." PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -v $test_dirs -m 'isolated_distributed'
+        PYTHONPATH="." PY_IGNORE_IMPORTMISMATCH=1 python3 -m pytest -s -v $test_dirs -m 'isolated_distributed'
     fi
 fi

imperative/python/test/unit/core/test_function.py

@@ -272,7 +272,7 @@ def test_none_in_out_grad():
     )
 
 
-def test_zero_grad():
+def test_clear_grad():
     class StopGradient(Function):
         def forward(self, a):
             return a

imperative/python/test/unit/core/test_indexing_op.py

@@ -556,18 +556,16 @@ def test_advance_indexing_with_bool(test_varnode):
     aa[bb] = False
     np.testing.assert_equal(a, aa.numpy())
 
-    # XXX: trace does not expect empty condtake tensor
-    if not use_symbolic_shape():
-        a = np.ones((2, 2), dtype=np.int32)
-        b = np.array([[False, False], [False, False]])
-        aa = Tensor(a)
-        bb = Tensor(b)
-        np.testing.assert_equal(a[b], aa[b].numpy())
-        np.testing.assert_equal(a[b], aa[bb].numpy())
-
-        b = np.array([False, False])
-        bb = Tensor(b)
-        np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape))  # FIXME
+    a = np.ones((2, 2), dtype=np.int32)
+    b = np.array([[False, False], [False, False]])
+    aa = Tensor(a)
+    bb = Tensor(b)
+    np.testing.assert_equal(a[b], aa[b].numpy())
+    np.testing.assert_equal(a[b], aa[bb].numpy())
+
+    b = np.array([False, False])
+    bb = Tensor(b)
+    np.testing.assert_equal(a[b], aa[bb].numpy().reshape(a[b].shape))
 
     a = np.arange(576).reshape(2, 3, 4, 3, 4, 2).astype("int32")
     aa = Tensor(a)

imperative/python/test/unit/jit/test_tracing.py

@@ -521,10 +521,11 @@ def test_trace_valid_broadcast():
     f(x2, shape)
 
 
-def test_clip():
+@pytest.mark.parametrize("trace_mode", [False, True])
+def test_clip(trace_mode):
     x = tensor(np.random.randn(10, 10))
 
-    @trace(symbolic=True)
+    @trace(symbolic=trace_mode)
     def f(x, lower, upper):
         y = F.clip(x, lower, upper)
         return y
@@ -532,6 +533,9 @@ def test_clip():
     for i in range(3):
         f(x, tensor([0]), tensor([1]))
 
+    for i in range(3):
+        f(x, tensor([5]), tensor([4]))
+
 
 # test returning noncontiguous tensor from trace
 def test_slice():