Function test_matmul_scale

imperative/python/test/unit/core/test_custom_op.py:174–183

(m=1, k=1, n=1, seed=2021)

        return out

    def test_matmul_scale(m=1, k=1, n=1, seed=2021):
        lhs_np, rhs_np, _, scale = gen_matmul_data(seed, m, k, n)
        custom_lhs, custom_rhs = Tensor(lhs_np), Tensor(rhs_np)
        builtin_lhs, builtin_rhs = Tensor(lhs_np), Tensor(rhs_np)

        custom_func = MatMulScale(scale=scale)
        custom_out = custom_func(custom_lhs, custom_rhs)
        builtin_out = builtin_func(builtin_lhs, builtin_rhs, scale)

        np.testing.assert_allclose(custom_out, builtin_out, rtol=1e-3, atol=1e-5)

    def test_matmul_scale_trace(m=1, k=1, n=1, seed=2021):
        @jit.trace(capture_as_const=True)