No public description

parallax/offload.py

@@ -104,7 +104,7 @@ def forward(model_chunk, intermediate_result):
   # example element 0 is the inputs to layer 0. The last element of the array
   # are the final outputs (which become the inputs to the loss function).
   saved_intermediates = [inputs]
-  for chunk in model.layers:
+  for chunk in model.layers:  # pyrefly: ignore[missing-attribute]
     logits = forward(chunk, saved_intermediates[-1])
     saved_intermediates.append(logits)
 
@@ -132,14 +132,14 @@ def offload_backward(
     A full nnx.State mapping representing a composite of the grads from all
     layers.
   """
-  if len(saved_intermediates) != len(model.layers) + 1:
+  if len(saved_intermediates) != len(model.layers) + 1:  # pyrefly: ignore[missing-attribute]
     raise ValueError(
         'The length of `saved_intermediates` must match the number of model '
         'layers plus one.'
     )
 
   layer_grads = {}
-  for i, chunk in reversed(list(enumerate(model.layers))):
+  for i, chunk in reversed(list(enumerate(model.layers))):  # pyrefly: ignore[missing-attribute]
     # pylint: disable=cell-var-from-loop
     chunk_input = saved_intermediates[i]
     graphdef, state = nnx.split(chunk)
@@ -189,13 +189,13 @@ def remat_model(model: nnx.Module) -> nnx.Module:
   """Takes an NNX Module and returns one with all layers rematerialized."""
   # TODO(jeffcarp): Generalize this to work with non-Sequential models.
   new_model = nnx.clone(model)
-  for i, layer in enumerate(new_model.layers):
+  for i, layer in enumerate(new_model.layers):  # pyrefly: ignore[missing-attribute]
     signature = inspect.signature(layer.__call__)
     unbound_call = layer.__class__.__call__
     # Shift static_argnums by 1 because of `self` in unbound call.
     static_argnums = tuple(i + 1 for i in range(1, len(signature.parameters)))
     rematted_call = nnx.remat(unbound_call, static_argnums=static_argnums)
-    new_model.layers[i] = RemattedLayer(layer, rematted_call)
+    new_model.layers[i] = RemattedLayer(layer, rematted_call)  # pyrefly: ignore[missing-attribute]
   return new_model
 
 

parallax/offload_test.py

@@ -100,8 +100,8 @@ def loss_fn(model, inputs):
     np.testing.assert_array_equal(act_loss, exp_loss)
     # Assert all model weights match after gradient update.
     np.testing.assert_allclose(
-        actual_model.layers[0].kernel[...],
-        reference_model.layers[0].kernel[...],
+        actual_model.layers[0].kernel[...],  # pyrefly: ignore[missing-attribute]
+        reference_model.layers[0].kernel[...],  # pyrefly: ignore[missing-attribute]
         atol=1e-3,
         rtol=1e-1,
     )

parallax/sharding/auto_shard.py

@@ -90,7 +90,7 @@ def get_shardings(
     for i in range(var.aval.ndim - 1, -1, -1):  # pytype: disable=attribute-error
       root = graph.get_root((var, i))
       if (
-          (model_axis is None)
+          (model_axis is None)  # pyrefly: ignore[unbound-name]
           or ((root, model_axis) in edges)
           or (var.aval.shape[i] < min_shard_size)
           or dim_sharded

parallax/sharding/base_test.py

@@ -53,7 +53,7 @@ def forward(params, inputs):
       return model.apply(params, inputs)
 
     output = forward(params, x)
-    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)
+    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)  # pyrefly: ignore[bad-argument-type]
 
     # Verify compiled shardings.
     compiled = forward.lower(params, x).compile()  # type: ignore
@@ -84,7 +84,7 @@ def forward(params, inputs):
       return model.apply(params, inputs)
 
     output = forward(params, x)
-    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)
+    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)  # pyrefly: ignore[bad-argument-type]
 
     # Verify compiled shardings.
     compiled = forward.lower(params, x).compile()  # type: ignore
@@ -115,7 +115,7 @@ def forward(params, inputs):
 
     sharded_forward = base.jit(forward, strategy=base.ShardingStrategy.DDP)
     output = sharded_forward(params, x)
-    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)
+    np.testing.assert_array_almost_equal(output, reference_output, decimal=6)  # pyrefly: ignore[bad-argument-type]
 
     # Verify compiled shardings.
     compiled = sharded_forward.lower(params, x).compile()  # type: ignore