modelverse/kernel/modelverse_jit/jit_primitives.py

"""Primitive functions that interface with the Modelverse JIT."""

import modelverse_kernel.primitives as primitive_functions

# pylint: disable=I0011,C0103

def get_jit_enabled(**kwargs):
    """Checks if the JIT is enabled."""
    result, = yield [("CNV", [kwargs['mvk'].jit.jit_enabled])]
    raise primitive_functions.PrimitiveFinished(result)

def set_jit_enabled(a, **kwargs):
    """Enables or disables the JIT."""
    val, result = yield [("RV", [a]), ("CN", [])]
    kwargs['mvk'].jit.set_jit_enabled(val)
    raise primitive_functions.PrimitiveFinished(result)