def _round(operand, *, rounding_method):
    if (rounding_method is lax.RoundingMethod.AWAY_FROM_ZERO):
        sign = _sign(operand)
        operand *= sign
        floor = tf.math.floor(operand)
        operand -= floor
        cond = tf.math.equal(operand, tf.constant(np.array(0.5), operand.dtype))
        return (sign * (tf.where(cond, tf.constant(np.array(1), operand.dtype), tf.math.round(operand)) + floor))
    else:
        return tf.math.round(operand)
def _round(operand, *, rounding_method):
    if (rounding_method is lax.RoundingMethod.AWAY_FROM_ZERO):
        sign = _sign(operand)
        operand *= sign
        floor = tf.math.floor(operand)
        operand -= floor
        cond = tf.math.equal(operand, tf.constant(np.array(0.5), operand.dtype))
        return (sign * (tf.where(cond, tf.constant(np.array(1), operand.dtype), tf.math.round(operand)) + floor))
    else:
        return tf.math.round(operand)