flax

Added kernel_metadata/bias_metadata args to nnx layers

#5074

Merged