Make fp8 ops use explicit broadcasting.

[levskaya]

Flax libraries use only explicit broadcasting and our unit tests enforce this, fix recent fp8 ops to follow this convention.

[codecov-commenter]

Codecov Report

Merging #3374 (2c2bd58) into main (242f84c) will increase coverage by 0.28%.
Report is 10 commits behind head on main.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #3374      +/-   ##
==========================================
+ Coverage   83.22%   83.51%   +0.28%     
==========================================
  Files          55       56       +1     
  Lines        6571     6678     +107     
==========================================
+ Hits         5469     5577     +108     
+ Misses       1102     1101       -1     

Files	Coverage Δ
flax/linen/__init__.py	100.00% <100.00%> (ø)
flax/linen/fp8_ops.py	100.00% <100.00%> (ø)
flax/linen/module.py	92.04% <100.00%> (ø)
flax/training/train_state.py	100.00% <100.00%> (ø)

... and 1 file with indirect coverage changes