Function attention_tensor 

pub fn attention_tensor(
    batch: DimExpr,
    heads: DimExpr,
    seq_len: DimExpr,
    head_dim: DimExpr,
) -> DependentType

Attention tensor (batch, heads, seq_len, head_dim).