guards

Strict Mode Guards Injection.

Attributes

STRICT_RANK_HELPER

Functions

apply_strict_guards(→ List[libcst.Arg])

Wraps args with rank assertions.

Module Contents

guards.STRICT_RANK_HELPER = Multiline-String

Show Value

"""
def _check_rank(x, rank):
    if hasattr(x, 'ndim'): n = x.ndim
    elif hasattr(x, 'shape'): n = len(x.shape)
    else: return x
    assert n == rank, f"Strict Guard: Expected rank {rank}, got {n}"
    return x
"""

guards.apply_strict_guards(rewriter: Any, norm_args: List[libcst.Arg], details: Dict[str, Any], target_impl: Dict[str, Any]) → List[libcst.Arg]

Wraps args with rank assertions.