def METHOD_NAME(
    pred,
    true_fun,
    false_fun,
    operand,
    rng_key=None,
    substitute_stack=None,
    enum=False,
    first_available_dim=None,
):
    if enum:
        # TODO: support enumeration. note that pred passed to lax.cond must be scalar
        # which means that even simple conditions like `x == 0` can get complicated if
        # x is an enumerated discrete random variable
        raise RuntimeError("The cond primitive does not currently support enumeration")
    if substitute_stack is None:
        substitute_stack = []
    wrapped_true_fun = wrap_fn(true_fun, substitute_stack)
    wrapped_false_fun = wrap_fn(false_fun, substitute_stack)
    wrapped_operand = device_put((rng_key, operand))
    return lax.cond(pred, wrapped_true_fun, wrapped_false_fun, wrapped_operand)