examples/research_projects/performer/modeling_flax_performer_utils.py/body
def _denominator(t_slice_shape, precision, unroll=1):
    def fwd(qs, ks):
        def body(p, qk):
            q, k = qk
            p += k
            x = jnp.einsum("...m,...m->...", q, p, precision=precision)
            return p, x

        p = jnp.zeros(t_slice_shape)
        p, R = lax.scan(body, p, (qs, ks), unroll=unroll)
        return R, (qs, ks, p)

    def bwd(qkp, R_ct):
        def body(carry, qkx):
            p, p_ct = carry
            q, k, x_ct = qkx
            q_ct = jnp.einsum("...,...m->...m", x_ct, p, precision=precision)
            p_ct += jnp.einsum("...,...m->...m", x_ct, q, precision=precision)
            k_ct = p_ct
            p -= k
            return (p, p_ct), (q_ct, k_ct)

        qs, ks, p = qkp
        _, (qs_ct, ks_ct) = lax.scan(body, (p, jnp.zeros_like(p)), (qs, ks, R_ct), reverse=True, unroll=unroll)
        return (qs_ct, ks_ct)

    @jax.custom_vjp
    def _denominator_impl(qs, ks):
        R, _ = fwd(qs, ks)
        return R

    _denominator_impl.defvjp(fwd, bwd)

    return _denominator_impl