Support gemma4 decoder under explicit sharding mode

[csgoogle]

Summary

• Enables shard_mode=explicit (needed for tensor parallelism) with the gemma4 decoder, which previously hit a hard ValueError (Decoder 'gemma4' is not supported with 'explicit' sharding).
• Gemma4DecoderLayer now threads out_sharding hints through its attention and MoE MLP calls, and uses a mesh/shard_mode-aware with_logical_constraint helper (maybe_shard_with_logical) instead of raw nn.with_logical_constraint, which only auto-resolves sharding mismatches under Auto mode, not Explicit.
• Adds "gemma4" to the set of decoders permitted under shard_mode=explicit in types.py.

Test plan

• Ran gemma4-26b RL (GRPO) training end-to-end on TPU v7x8 with rollout_tensor_parallelism=4, shard_mode=explicit, including a captured profiler trace and checkpoint save.
• CI

[gemini-code-assist]

Code Review

This pull request enables explicit sharding support for the Gemma4 model by adding "gemma4" to the supported decoders list and refactoring gemma4.py to use custom logical constraints (maybe_shard_with_logical) instead of relying on flax.linen.with_logical_constraint. The reviewer feedback recommends a more robust approach for retrieving logical axis rules: instead of calling get_logical_axis_rules(), which relies on thread-local context, the code should directly reference the configured rules via config.logical_axis_rules across the newly added sharding calls.

[gemini-code-assist]

Using config.logical_axis_rules directly is more robust than relying on get_logical_axis_rules(), which reads from the thread-local context. During initialization (__init__), the thread-local context rules might not be fully established or could differ from the configured rules (e.g., in unit tests or setup phases). If you apply this change across all occurrences, you can also remove the import of get_logical_axis_rules.

Suggested change

	self.out_sharding = (
	create_sharding(self.mesh, self.activation_axis_names, rules=get_logical_axis_rules()) if self.mesh else None
	)
	self.out_sharding = (
	create_sharding(self.mesh, self.activation_axis_names, rules=config.logical_axis_rules) if self.mesh else None
	)

[gemini-code-assist]

Passing self.config.logical_axis_rules directly is more robust and consistent than relying on get_logical_axis_rules().

Suggested change

	return maybe_shard_with_logical(
	x,
	logical_axes=self.activation_axis_names,
	mesh=self.mesh,
	shard_mode=getattr(self.config, "shard_mode", ShardMode.AUTO),
	debug_sharding=getattr(self.config, "debug_sharding", False),
	extra_stack_level=1,
	rules=get_logical_axis_rules(),
	)
	return maybe_shard_with_logical(
	x,
	logical_axes=self.activation_axis_names,
	mesh=self.mesh,
	shard_mode=getattr(self.config, "shard_mode", ShardMode.AUTO),
	debug_sharding=getattr(self.config, "debug_sharding", False),
	extra_stack_level=1,
	rules=self.config.logical_axis_rules,
	)

[gemini-code-assist]

Use cfg.logical_axis_rules directly to ensure robustness and consistency with the other sharding constraint calls.

Suggested change

	inputs = maybe_shard_with_logical(
	inputs,
	logical_axes=("activation_batch", "activation_norm_length", "activation_embed"),
	mesh=self.mesh,
	shard_mode=getattr(self.config, "shard_mode", ShardMode.AUTO),
	debug_sharding=getattr(self.config, "debug_sharding", False),
	rules=get_logical_axis_rules(),
	)
	inputs = maybe_shard_with_logical(
	inputs,
	logical_axes=("activation_batch", "activation_norm_length", "activation_embed"),
	mesh=self.mesh,
	shard_mode=getattr(self.config, "shard_mode", ShardMode.AUTO),
	debug_sharding=getattr(self.config, "debug_sharding", False),
	rules=cfg.logical_axis_rules,
	)