RF test_compare_bc, test_logical_or

Author: albertz

tests/test_rf_math.py

@@ -3,7 +3,13 @@
 """
 
 from __future__ import annotations
+
 import _setup_test_env  # noqa
+
+import sys
+import unittest
+
+from returnn.util import better_exchook
 import returnn.frontend as rf
 from returnn.tensor import Tensor, Dim, TensorDict, batch_dim
 from rf_utils import run_model
@@ -30,6 +36,39 @@ def _forward_step(*, model: _Net, extern_data: TensorDict):
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_compare_bc():
+    beam_dim = Dim(3, name="beam")
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict({"idx": Tensor("idx", [batch_dim, beam_dim], dtype="int32", sparse_dim=in_dim)})
+
+    # noinspection PyShadowingNames,PyUnusedLocal
+    def _forward_step(*, model: rf.Module, extern_data: TensorDict):
+        idx = extern_data["idx"]
+        cond = rf.compare_bc(idx, "!=", rf.range_over_dim(in_dim))
+        cond.mark_as_default_output(shape=(batch_dim, beam_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+def test_logical_or():
+    beam_dim = Dim(3, name="beam")
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "a": Tensor("a", [batch_dim, beam_dim], dtype="bool"),
+            "b": Tensor("b", [batch_dim, beam_dim, in_dim], dtype="bool"),
+        }
+    )
+
+    # noinspection PyShadowingNames,PyUnusedLocal
+    def _forward_step(*, model: rf.Module, extern_data: TensorDict):
+        a, b = extern_data["a"], extern_data["b"]
+        cond = a | b
+        cond.mark_as_default_output(shape=(batch_dim, beam_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
 def test_squared_difference():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")
@@ -125,3 +164,26 @@ def _forward_step(*, model: rf.Module, extern_data: TensorDict):
         out.mark_as_default_output(shape=(batch_dim, time_dim))
 
     run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
+if __name__ == "__main__":
+    better_exchook.install()
+    if len(sys.argv) <= 1:
+        for k, v in sorted(globals().items()):
+            if k.startswith("test_"):
+                print("-" * 40)
+                print("Executing: %s" % k)
+                try:
+                    v()
+                except unittest.SkipTest as exc:
+                    print("SkipTest:", exc)
+                print("-" * 40)
+        print("Finished all tests.")
+    else:
+        assert len(sys.argv) >= 2
+        for arg in sys.argv[1:]:
+            print("Executing: %s" % arg)
+            if arg in globals():
+                globals()[arg]()  # assume function and execute
+            else:
+                eval(arg)  # assume Python code and execute