Fix linter

Signed-off-by: Vladimir Suvorov <suvorovv@google.com>

Author: mydatascience

tests/pipeline_parallelism_test.py

@@ -71,7 +71,15 @@ def assert_pipeline_same_output_and_grad(self, config, single_pipeline_stage_cla
         rngs = nnx.Rngs(params=0)
         single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode, rngs=rngs)
       else:
-        single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode)
+        try:
+          single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode)
+        except TypeError as exc:
+          if "rngs" not in str(exc):
+            raise
+          rngs = nnx.Rngs(params=0)
+          single_pipeline_stage = single_pipeline_stage_class(
+              config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
+          )
 
     def get_inputs(batch_size, sequence, features):
       """Get random inputs, and random dummy targets