Fix DeepSeek4 shape mismatch in mhc.py and moe.py

[olufiyin19]

Fix DeepSeek4 shape mismatch in mhc.py and moe.py

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request refactors DeepSeek4HyperHead in mhc.py to inherit from nnx.Module with custom parameters, updates the vocab size lookup for DeepSeek V4 Hash Routing in moe.py, and integrates deepstack_process in nnx_decoders.py while cleaning up duplicate code. Feedback highlights an unreachable duplicate return statement in DeepSeek4HyperHead.call and a duplicated comment line in moe.py.

[gemini-code-assist]

There is an unreachable duplicate return statement at the end of the __call__ method. The second return statement should be removed.

Suggested change

	return x
	return jnp.sum(x, axis=2)
	return x

[gemini-code-assist]

The comment # DeepSeek V4 Hash Routing is duplicated. Please remove the duplicate line.

Suggested change

	# DeepSeek V4 Hash Routing
	# DeepSeek V4 Hash Routing
	# DeepSeek V4 Hash Routing