Convert DenseGeneral to NNX

[bvandermoon]

Description

Finishing out PR 1604 for @cgarciae due to a lot of issues being faced during testing (we are both listed on the commit, but it doesn't look like using multiple PR authors is supported).

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
• Will be running a logits check for gpt3 which is the first model we are using

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.