    def test_constant_folding_mixed_usage(self):
        """
        Test a mix of all the ways to use constants:
        - Unbound constant that is not folded
        - Unbound constant folding input
        - Bound constant folding input
        - Bound constant that is not folded
        """
        dtype = "float16"

        N, K = IntImm(13), IntImm(33)
        input_0 = Tensor(shape=[N, N], dtype=dtype, name="input_0", is_input=True)

        # Unbound, unfolded constant
        w1_ait = Tensor(shape=[N, N], dtype=dtype, name="w1")

        x1_ait = ops.elementwise(FuncEnum.MUL)(w1_ait, input_0)

        # Unbound folded constant
        w2_ait = Tensor(shape=[N, K], dtype=dtype, name="w2")

        # Bound folded constants
        w3_ait = Tensor(shape=[N, K], dtype=dtype, name="w3")
        w4_ait = Tensor(shape=[N, K], dtype=dtype, name="w4")

        x2_ait = ops.elementwise(FuncEnum.MUL)(w2_ait, w3_ait)
        x3_ait = ops.gemm_rcr()(x2_ait, w4_ait)

        x4_ait = ops.elementwise(FuncEnum.MUL)(x3_ait, x1_ait)

        # Bound unfolded constant
        w5_ait = Tensor(shape=[N, N], dtype=dtype, name="w5")
        output = ops.elementwise(FuncEnum.MUL)(w5_ait, x4_ait)
        output._attrs["is_output"] = True
        output._attrs["name"] = "output"

        input_pt = get_random_torch_tensor((N.value(), N.value()), dtype)
        w1_pt = get_random_torch_tensor((N.value(), N.value()), dtype)
        w2_pt = get_random_torch_tensor((N.value(), K.value()), dtype)
        w3_pt = get_random_torch_tensor((N.value(), K.value()), dtype)
        w4_pt = get_random_torch_tensor((N.value(), K.value()), dtype)
        w5_pt = get_random_torch_tensor((N.value(), N.value()), dtype)

        x1_pt = w1_pt * input_pt
        x2_pt = w2_pt * w3_pt
        x3_pt = torch.nn.functional.linear(x2_pt, w4_pt)
        x4_pt = x3_pt * x1_pt
        output_pt = w5_pt * x4_pt

        mod = compile_model(
            output,
            detect_target(),
            "./tmp",
            "test_constant_folding_mixed_usage",
            constants={"w3": w3_pt, "w4": w4_pt, "w5": w5_pt},
        )

        self.assertSetEqual(
            set(mod.get_constant_folding_input_names()),
            # This is not the only input, but it's the only one we can set.
            {"w2"},
        )

        self.assertSetEqual(set(mod.get_constant_names()), {"w1", "w2"})

        output = torch.empty_like(output_pt)
        # Unset constant W2
        self.assertRaises(RuntimeError, mod.run_with_tensors, [input_pt], [output])
        mod.set_constant_with_tensor("w2", w2_pt)
        # Unset constant W1
        self.assertRaises(RuntimeError, mod.run_with_tensors, [input_pt], [output])
        mod.set_constant_with_tensor("w1", w1_pt)

        mod.run_with_tensors([input_pt], [output])
        torch.testing.assert_close(output, output_pt, atol=1e-1, rtol=1e-1)
Test a mix of all the ways to use constants:
- Unbound constant that is not folded
- Unbound constant folding input
- Bound constant folding input
- Bound constant that is not folded
tests/unittest/compiler/test_constant_folding.py