flax

c6036553 - Added kernel_metadata/bias_metadata args to nnx layers

Added kernel_metadata/bias_metadata args to nnx layers Description: Added kernel_metadata/bias_metadata args to nnx layers: - [x] nnx.Linear - [x] nnx.LinearGeneral - [ ] nnx.Conv - [ ] ...