Support dot broadcast

[prbzrg]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

We need this patch to be able to use something like:

DifferentiationInterface.jacobian.(x -> Lux.apply(nns, x, ps), AutoZygote(), eachcol(r))