lint(grpo): silence pylint false positive on TrainStateNNX kwarg

pylint resolves train_state_nnx via the stale non-editable site-packages copy where TrainStateNNX.__init__ predates reference_model, so the keyword call in init_state_fn trips E1123. The local class at src/maxtext/common/train_state_nnx.py does accept reference_model. Disable the rule on this single call.

Author: ecnal-cienet

src/maxtext/experimental/rl/grpo_trainer.py

@@ -798,6 +798,7 @@ def init_state_fn():
         optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
         # Reference uses the same init seed so it starts identical to the policy.
         reference_model = _create_model_partial()
+        # pylint: disable-next=unexpected-keyword-arg
         return train_state_nnx.TrainStateNNX(nnx_model, optimizer, reference_model=reference_model)
 
     else: