Method forward

axlearn/common/attention_test.py:4525–4537

(self, data, **layer_kwargs)

        # [batch, length, dim].
        output = self.stack(data, **layer_kwargs)
        x = output.data
        x_mean = jnp.mean(x, axis=1, keepdims=True)
        # [batch, length].
        x_var = jnp.sum((x - x_mean) ** 2, axis=-1)
        loss = jnp.mean(x_var)
        if cfg.output_self_attention_kv_state:
            return loss, {"mean": x_mean, "self_attention_kv_state": output.self_attention_kv_state}
        return loss, {"mean": x_mean}


def _recursive_stack(inputs: Nested[Tensor], axis=0):
    def stack(*xs):
        return jnp.stack(xs, axis=axis)

    return {"layer": utils.vectorized_tree_map(stack, *inputs.values())}