Temporarily disable channel-wise pruning in praxis to align its imple…

…mentation with the sparsity/jax channelwise pruning implementation.

PiperOrigin-RevId: 629473113

praxis/layers/quantization/sparsity/sparsifier.py

@@ -542,6 +542,10 @@ def sparsifiy(
     ):
       return weight
 
+    assert (
+        self.sparsity.sparsity_type != SparsityType.CHANNELWISE_PRUNING
+    ), 'Channel-wise pruning is temporarily disabled'
+
     step = self.get_var('step')
 
     # Return without updating mask if in we want to do mixed sparsity for

praxis/layers/quantization/sparsity/sparsifier_test.py

@@ -16,6 +16,8 @@
 """Tests for sparse_base_layer."""
 
 import copy
+import unittest
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -30,6 +32,7 @@
 from praxis.layers.quantization.sparsity import sparsity_hparams
 from praxis.layers.quantization.sparsity import sparsity_modes
 
+
 instantiate = base_layer.instantiate
 NON_TRAINABLE = base_layer.NON_TRAINABLE
 PARAMS = base_layer.PARAMS
@@ -1222,6 +1225,7 @@ def update(test_layer, params, inputs, updated_weights):
           ),
       )
 
+  @unittest.skip('Channel-wise pruning is temporarily disabled')
   @parameterized.named_parameters(
       ('column_wise', -1),
       ('row_wise', -2),