Function _args

tests/test_operator_inputs.py:14–30

(**overrides)

from rl_engine.kernels.gtest.operator_specs import (
    make_candidate,
    make_operator_case,
    operator_names,
)


def _args(**overrides):
    values = {
        "batch": 1,
        "seq": 2,
        "vocab": 17,
        "seed": 123,
        "input_mode": "constant",
        "constant_value": 0.5,
        "token_value": 3,
        "normalized_dim": 128,
        "n_heads": 32,
        "head_dim": 128,
        "k_dim": 16,
        "n_dim": 32,
        "theta": 1.0e6,