pypto.cond

Applicable Products

• Ascend 950PR/Ascend 950DT: Supported
• Atlas A3 training products/Atlas A3 inference products: Supported
• Atlas A2 training products/Atlas A2 inference products: Supported

Description

Defines an if conditional operation to implement the if functionality in Python.

Prototype

cond(scalar: SymInt) -> pypto_impl.RecordIfBranch

Parameters

Return Value

pypto_impl.RecordIfBranch: Returns a conditional branch object for use in Python if statements.

Constraints

• Must be used together with Python if, elif, and else statements.
• The conditional expression is recorded into the computation graph.
• Nested conditional statements are supported.
• When the function is not decorated with @pypto.frontend.jit or @pypto.frontend.function, the conditional expression must be wrapped with pypto.cond.

Example

# Without a decorator, wrap the conditional expression with pypto.cond.
def kernel():
    ...
    for s2_idx in pypto.loop(0, 10, 1, power_of_2(max_unroll_times), name="LOOP_L0_bIdx_mla_prolog", idx_name="b_idx"):
        if pypto.cond(pypto.is_loop_end(s2_idx, bn_per_batch)):
            ...

# With a decorator, no pypto.cond wrapping is required.
@pypto.frontend.jit
def kernel():
    ...
    for s2_idx in pypto.loop(0, 10, 1, power_of_2(max_unroll_times), name="LOOP_L0_bIdx_mla_prolog", idx_name="b_idx"):
        if pypto.is_loop_end(s2_idx, bn_per_batch):
            ...