Support for `allocate(backend, desired_axes)` in addition to `allocate(backend, desired_size)`

Hey KA.jl team, thanks for this amazing package! I'm very impressed at how uniquely easy it is to write portable GPU code using KA.jl, and I wish other programming ecosystems had this awesome power.

Recently, I've been rewriting a lot of my Julia code to use axes instead of size, and I noticed that allocate(backend, T, desired_axes) doesn't work, while allocate(backend, T, desired_size) does. This would be a nice feature to bring the allocate interface in line with Base.similar, which accepts size or axes. Are there any plans to support this? I would be happy to help implement this myself if someone can point me to the right code paths to looks at.

Note that this would be helpful for API compatibility even if KA.jl doesn't plan to support offset axes or symbolic axes -- we could just throw on any axis that isn't a 1:n or Base.OneTo(n).

Start at the allocate(backend, T, desired_size) entry point and compare its interface with Base.similar. Done means allocate also accepts desired_axes for standard 1:n or Base.OneTo(n) axes, rejects unsupported axes as described, and has coverage for the new behavior.