    def _test_1_shape(
        self,
        batch_sizes,
        ms,
        ns,
        ks,
        test_name,
        expected_max_read_t,
        expected_read_types,
        expected_op_t,
        expected_data_t,
        dtype="float16",
    ):
        """
        Tests tanh(A(B, 1, 1, M, K, 1) + B(N, N, 1, K, M)).
        """

        batch_dim = shape_utils.gen_int_var_min_max(batch_sizes)
        m_dim = shape_utils.gen_int_var_min_max(ms)
        n_dim = shape_utils.gen_int_var_min_max(ns)
        k_dim = shape_utils.gen_int_var_min_max(ks)

        X1 = Tensor(
            shape=[batch_dim, IntImm(1), IntImm(1), m_dim, k_dim, IntImm(1)],
            dtype=dtype,
            name="input0",
            is_input=True,
        )
        X2 = Tensor(
            shape=[n_dim, n_dim, IntImm(1), k_dim, m_dim],
            dtype=dtype,
            name="input1",
            is_input=True,
        )
        X3 = ops.elementwise(FuncEnum.ADD)(X1, X2)
        X4 = ops.elementwise(FuncEnum.TANH)(X3)
        X4._attrs["name"] = "output0"
        X4._attrs["is_output"] = True
        self.assertEqual(
            X4._attrs["shape"], [batch_dim, n_dim, n_dim, m_dim, k_dim, m_dim]
        )

        target = detect_target()
        module = compile_model(X4, target, "./tmp", test_name)

        debug_sorted_graph = module.debug_sorted_graph
        sorted_ops = graph_utils.get_sorted_ops(debug_sorted_graph)
        self.assertEqual(len(sorted_ops), 1)
        self.assertEqual(sorted_ops[0]._attrs["op"], "fused_elementwise")
        self.assertEqual(sorted_ops[0]._attrs["max_read_t"], expected_max_read_t)
        self.assertEqual(
            self._get_sorted_read_types(sorted_ops[0]), expected_read_types
        )
        self.assertEqual(sorted_ops[0]._attrs["op_t"], expected_op_t)
        self.assertEqual(sorted_ops[0]._attrs["data_t"], expected_data_t)

        for batch_size, m, n, k in itertools.product(batch_sizes, ms, ns, ks):
            x1_pt = get_random_torch_tensor([batch_size, 1, 1, m, k, 1], dtype=dtype)
            x2_pt = get_random_torch_tensor([n, n, 1, k, m], dtype=dtype)
            x4_pt = torch.tanh(x1_pt + x2_pt)
            inputs = {"input0": x1_pt, "input1": x2_pt}
            x4 = torch.empty_like(x4_pt)
            module.run_with_tensors(inputs, [x4])
            self.assertTrue(torch.allclose(x4, x4_pt, atol=1e-2, rtol=1e-2))
Tests tanh(A(B, 1, 1, M, K, 1) + B(N, N, 1, K, M)).
tests/unittest/ops/test_fused_elementwise_broadcast.py