Convert DenseGeneral to NNX

[bvandermoon]

Description

Finishing out PR 1604 for @cgarciae

Note: I updated the logits in golden_data_grpo_default.jsonl to the NNX values. The NNX logits don't match the Linen ones since the RNG keys are setup differently.

Previous PR description:

This commit converts DenseGeneral to NNX and creates a dense_general to interface with it through a Linen wrapper. dense_general contains all the same arguments as the Linen version but adds two additional ones:

• input_shape: the expected shape of the input.
• in_features: an int or tuple representing the input features.

Only one of them can be set at a time.

Tests

• Unit tests
• We may end up needing to regenerate and update other golden logits. The logits are not the same for NNX/Linen due to the way the RNG keys are handled and the weights are initialized

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.