def _abs_taylor_rule(x, series_in, **params): x, = x primal_out = lax.abs_p.bind(x, **params) negs = lax.select(lax.lt(x, 0.0), lax.full_like(x, -1), lax.full_like(x, 1.0)) fix_sign = lambda y: negs * y series_out = [fix_sign(*terms_in, **params) for terms_in in zip(*series_in)] return primal_out, series_out
def while_cond_fun(loop_carry):
i, _ = loop_carry
return lax.lt(i, upper)
def _fori_cond_fun(loop_carry): i, upper, _ = loop_carry return lax.lt(i, upper)
