flax
15d6857e - [linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.

Commit

2 years ago

[linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case. Added test for scalar stats, which were broken by previous change. PiperOrigin-RevId: 549263368

Author

chr1sj0nes

Committer

a-googler

Parents

881c4495

flax 15d6857e - [linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.

flax
15d6857e - [linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.