[TK] Add support for ops required for Flash Attention 2

[Groverkss]

Add new ops:

• tkl.exp2 (math)
• tkl.max (reduce max)
• tkl.sum (reduce sum)
• tkl.broadcast (broadcast leading dims)
• tkl.broadcast_in_dim (broadcast specific dimensions)
• tkl.transpose (transpose)

[stellaraccident]

Ok to land these as-is since prototyping. But as we've noted on Discord, if we're wanting to go much further, we need to step back and work some of this more systematically. But for now, onward...