Function get_numpy_attention_input

tests/python/relax/test_codegen_cutlass.py:1006–1014

(q_shape, k_shape, v_shape, bias_shape, dtype)

def get_numpy_attention_input(q_shape, k_shape, v_shape, bias_shape, dtype):
    q = np.random.randn(*q_shape).astype(dtype)
    k = np.random.randn(*k_shape).astype(dtype)
    v = np.random.randn(*v_shape).astype(dtype)
    if not bias_shape == "none":
        bias = np.random.randn(*bias_shape).astype(dtype)
    else:
        bias = None
    return q, k, v, bias


def test_attention_rewrite_offload(attention_rewrite_size):