Fixed op tests

tests/jax/graphs/test_example_graph.py

@@ -16,7 +16,6 @@ def example_graph(x: jax.Array, y: jax.Array) -> jax.Array:
 
 
 @pytest.mark.push
-@pytest.mark.nightly
 @pytest.mark.parametrize(
     ["x_shape", "y_shape"],
     [

tests/jax/multichip/manual/all_gather.py

@@ -6,10 +6,9 @@
 import jax.numpy as jnp
 import pytest
 from infra import make_partition_spec, run_multichip_test_with_random_inputs
-from utils import TestCategory, compile_fail
+from utils import compile_fail
 
 
-@pytest.mark.record_test_properties(test_category=TestCategory.MULTICHIP_TEST.value)
 @pytest.mark.parametrize(
     ("x_shape", "mesh_shape", "axis_names"), [((8192, 784), (2,), ("batch",))]
 )

tests/jax/multichip/manual/data_paralelism.py

@@ -6,10 +6,9 @@
 import jax.numpy as jnp
 import pytest
 from infra import make_partition_spec, run_multichip_test_with_random_inputs
-from utils import TestCategory, compile_fail
+from utils import compile_fail
 
 
-@pytest.mark.record_test_properties(test_category=TestCategory.MULTICHIP_TEST.value)
 @pytest.mark.parametrize(
     [
         "batch_shape",

tests/jax/multichip/manual/psum.py

@@ -6,10 +6,9 @@
 import jax.numpy as jnp
 import pytest
 from infra import make_partition_spec, run_multichip_test_with_random_inputs
-from utils import TestCategory, compile_fail
+from utils import compile_fail
 
 
-@pytest.mark.record_test_properties(test_category=TestCategory.MULTICHIP_TEST.value)
 @pytest.mark.parametrize(
     ["batch_shape", "W1_shape", "B1_shape", "mesh_shape", "axis_names"],
     [

tests/jax/multichip/manual/psum_scatter.py

@@ -6,10 +6,9 @@
 import jax.numpy as jnp
 import pytest
 from infra import make_partition_spec, run_multichip_test_with_random_inputs
-from utils import TestCategory, compile_fail
+from utils import compile_fail
 
 
-@pytest.mark.record_test_properties(test_category=TestCategory.MULTICHIP_TEST.value)
 @pytest.mark.parametrize(
     ["batch_shape", "W1_shape", "B1_shape", "mesh_shape", "axis_names"],
     [

tests/jax/multichip/manual/unary_eltwise.py

@@ -6,10 +6,9 @@
 import jax.numpy as jnp
 import pytest
 from infra import make_partition_spec, run_multichip_test_with_random_inputs
-from utils import TestCategory, compile_fail
+from utils import compile_fail
 
 
-@pytest.mark.record_test_properties(test_category=TestCategory.MULTICHIP_TEST.value)
 @pytest.mark.parametrize(
     ("x_shape", "mesh_shape", "axis_names"), [((256, 256), (1, 2), ("x", "y"))]
 )

tests/jax/ops/test_and.py

@@ -1,5 +1,34 @@
-# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
 #
 # SPDX-License-Identifier: Apache-2.0
 
-# TODO add tests for `stablehlo.and`.
+import jax
+import jax.numpy as jnp
+import pytest
+from infra import random_tensor, run_op_test
+from utils import TestCategory, convert_output_to_bfloat16
+
+
+@pytest.mark.push
+@pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    test_category=TestCategory.OP_TEST.value,
+    jax_op_name="jax.numpy.logical_and",
+    shlo_op_name="stablehlo.and{LOGICAL}",
+)
+@pytest.mark.parametrize(
+    ["shape"],
+    [
+        [(32, 32)],
+        [(64, 64)],
+    ],
+    ids=lambda val: f"{val}",
+)
+def test_logical_and(shape: tuple):
+    @convert_output_to_bfloat16
+    def logical_and(a: jax.Array, b: jax.Array) -> jax.Array:
+        return jnp.logical_and(a, b)
+
+    lhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=3)
+    rhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=6)
+    run_op_test(logical_and, [lhs, rhs])

tests/jax/ops/test_not.py

@@ -1,5 +1,33 @@
-# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
 #
 # SPDX-License-Identifier: Apache-2.0
 
-# TODO add tests for `stablehlo.not`.
+import jax
+import jax.numpy as jnp
+import pytest
+from infra import random_tensor, run_op_test
+from utils import TestCategory, convert_output_to_bfloat16
+
+
+@pytest.mark.push
+@pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    test_category=TestCategory.OP_TEST.value,
+    jax_op_name="jax.numpy.logical_not",
+    shlo_op_name="stablehlo.not{LOGICAL}",
+)
+@pytest.mark.parametrize(
+    ["shape"],
+    [
+        [(32, 32)],
+        [(64, 64)],
+    ],
+    ids=lambda val: f"{val}",
+)
+def test_logical_not(shape: tuple):
+    @convert_output_to_bfloat16
+    def logical_not(a: jax.Array) -> jax.Array:
+        return jnp.logical_not(a)
+
+    input = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=3)
+    run_op_test(logical_not, [input])

tests/jax/ops/test_or.py

@@ -1,5 +1,34 @@
-# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
 #
 # SPDX-License-Identifier: Apache-2.0
 
-# TODO add tests for `stablehlo.or`.
+import jax
+import jax.numpy as jnp
+import pytest
+from infra import random_tensor, run_op_test
+from utils import TestCategory, convert_output_to_bfloat16
+
+
+@pytest.mark.push
+@pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    test_category=TestCategory.OP_TEST.value,
+    jax_op_name="jax.numpy.logical_or",
+    shlo_op_name="stablehlo.or{LOGICAL}",
+)
+@pytest.mark.parametrize(
+    ["shape"],
+    [
+        [(32, 32)],
+        [(64, 64)],
+    ],
+    ids=lambda val: f"{val}",
+)
+def test_logical_or(shape: tuple):
+    @convert_output_to_bfloat16
+    def logical_or(a: jax.Array, b: jax.Array) -> jax.Array:
+        return jnp.logical_or(a, b)
+
+    lhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=3)
+    rhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=6)
+    run_op_test(logical_or, [lhs, rhs])

tests/jax/ops/test_tanh.py

@@ -8,19 +8,19 @@
 import jax.numpy as jnp
 import pytest
 from infra import run_op_test_with_random_inputs
-from utils import record_unary_op_test_properties
+from utils import TestCategory
 
 
 @pytest.mark.push
 @pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    test_category=TestCategory.OP_TEST.value,
+    jax_op_name="jax.numpy.subtract",
+    shlo_op_name="stablehlo.subtract",
+)
 @pytest.mark.parametrize("x_shape", [(32, 32), (64, 64)], ids=lambda val: f"{val}")
-def test_tanh(x_shape: tuple, record_tt_xla_property: Callable):
+def test_tanh(x_shape: tuple):
     def tanh(x: jax.Array) -> jax.Array:
         return jnp.tanh(x)
 
-    record_unary_op_test_properties(
-        record_tt_xla_property,
-        "jax.numpy.tanh",
-        "stablehlo.tanh",
-    )
     run_op_test_with_random_inputs(tanh, [x_shape])

tests/jax/ops/test_xor.py

@@ -1,5 +1,34 @@
-# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
 #
 # SPDX-License-Identifier: Apache-2.0
 
-# TODO add tests for `stablehlo.xor`.
+import jax
+import jax.numpy as jnp
+import pytest
+from infra import random_tensor, run_op_test
+from utils import TestCategory, convert_output_to_bfloat16
+
+
+@pytest.mark.push
+@pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    test_category=TestCategory.OP_TEST.value,
+    jax_op_name="jax.numpy.logical_xor",
+    shlo_op_name="stablehlo.xor{LOGICAL}",
+)
+@pytest.mark.parametrize(
+    ["shape"],
+    [
+        [(32, 32)],
+        [(64, 64)],
+    ],
+    ids=lambda val: f"{val}",
+)
+def test_logical_xor(shape: tuple):
+    @convert_output_to_bfloat16
+    def logical_xor(a: jax.Array, b: jax.Array) -> jax.Array:
+        return jnp.logical_xor(a, b)
+
+    lhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=3)
+    rhs = random_tensor(shape, jnp.int32, minval=0, maxval=2, random_seed=6)
+    run_op_test(logical_xor, [lhs, rhs])

tests/utils.py

@@ -15,7 +15,6 @@ class TestCategory(Enum):
     OP_TEST = "op_test"
     GRAPH_TEST = "graph_test"
     MODEL_TEST = "model_test"
-    MULTICHIP_TEST = "multichip_test"
     OTHER = "other"