examples/research_projects/performer/modeling_flax_performer_utils.py/_numerator_impl
def _numerator(z_slice_shape, precision, unroll=1):
    def fwd(qs, ks, vs):
        def body(p, qkv):
            (q, k, v) = qkv
            p += jnp.einsum("...m,...d->...md", k, v, precision=precision)
            X_slice = jnp.einsum("...m,...md->...d", q, p, precision=precision)
            return p, X_slice

        init_value = jnp.zeros(z_slice_shape)
        p, W = lax.scan(body, init_value, (qs, ks, vs), unroll=unroll)
        return W, (p, qs, ks, vs)

    def bwd(pqkv, W_ct):
        def body(carry, qkv_xct):
            p, p_ct = carry
            q, k, v, x_ct = qkv_xct
            q_ct = jnp.einsum("...d,...md->...m", x_ct, p, precision=precision)
            p_ct += jnp.einsum("...d,...m->...md", x_ct, q, precision=precision)
            k_ct = jnp.einsum("...md,...d->...m", p_ct, v, precision=precision)
            v_ct = jnp.einsum("...md,...m->...d", p_ct, k, precision=precision)
            p -= jnp.einsum("...m,...d->...md", k, v, precision=precision)
            return (p, p_ct), (q_ct, k_ct, v_ct)

        p, qs, ks, vs = pqkv
        _, (qs_ct, ks_ct, vs_ct) = lax.scan(
            body, (p, jnp.zeros_like(p)), (qs, ks, vs, W_ct), reverse=True, unroll=unroll
        )
        return qs_ct, ks_ct, vs_ct

    @jax.custom_vjp
    def _numerator_impl(qs, ks, vs):
        W, _ = fwd(qs, ks, vs)
        return W

    _numerator_impl.defvjp(fwd, bwd)

    return _numerator_impl