Rename create_aux_variables to create_sparsity_variables for readabil…

…ity. PiperOrigin-RevId: 605141468
google · Feb 8, 2024 · bfb6127 · bfb6127
1 parent 775727f
commit bfb6127
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/praxis/layers/quantization/attentions.py b/praxis/layers/quantization/attentions.py
@@ -87,7 +87,7 @@ def setup(self) -> None:
         weight_params=pc,
         scale_shape=scale_shape,
     )
-    self.create_aux_variables('w', pc, scale_shape=scale_shape)
+    self.create_sparsity_variables('w', pc, scale_shape=scale_shape)
 
     if self.use_bias:
       if self.is_output_projection:
@@ -335,7 +335,7 @@ def setup(self) -> None:
         weight_params=pc,
         scale_shape=[3] + hd_shape,
     )
-    self.create_aux_variables('w', pc, scale_shape=[3] + hd_shape)
+    self.create_sparsity_variables('w', pc, scale_shape=[3] + hd_shape)
     if self.use_bias:
       # Combined bias weight for q, k, v projections.
       pc_bias = WeightHParams(

diff --git a/praxis/layers/quantization/linears.py b/praxis/layers/quantization/linears.py
@@ -147,7 +147,7 @@ def setup(self) -> None:
           weight_params=wp_a,
           scale_shape=[self.rank],
       )
-      self.create_aux_variables(
+      self.create_sparsity_variables(
           'w_a',
           wp_a,
           scale_shape=[self.rank],
@@ -162,7 +162,7 @@ def setup(self) -> None:
           weight_params=wp_b,
           scale_shape=[self.output_dims],
       )
-      self.create_aux_variables(
+      self.create_sparsity_variables(
           'w_b',
           wp_b,
           scale_shape=[self.output_dims],
@@ -185,7 +185,7 @@ def setup(self) -> None:
           weight_params=weight_hparams,
           scale_hparams=scale_hparams,
       )
-      self.create_aux_variables(
+      self.create_sparsity_variables(
           'w',
           weight_hparams,
           scale_shape=[self.output_dims],

diff --git a/praxis/layers/quantization/sparsity/sparsifier.py b/praxis/layers/quantization/sparsity/sparsifier.py
@@ -54,7 +54,7 @@ def sr_ste(
 
   Algorithm description: https://arxiv.org/abs/2102.04010
 
-  The last arguement is forced to be static to simplify
+  The last argument is forced to be static to simplify
     the implementation.
 
   Args:
@@ -129,7 +129,7 @@ class SparsityBaseLayer(base_layer.BaseLayer):
 
   sparsity: Optional[SparsityHParams] = None
 
-  def create_aux_variables(
+  def create_sparsity_variables(
       self,
       name: str,
       weight_hparams: WeightHParams,

diff --git a/praxis/layers/quantization/sparsity/sparsifier_test.py b/praxis/layers/quantization/sparsity/sparsifier_test.py
@@ -60,7 +60,7 @@ def setup(self):
     name = 'w'
     self.create_variable(name, weight_hp)
     self.create_child('einsum', self.einsum_tpl.clone())
-    self.create_aux_variables(name, weight_hp)
+    self.create_sparsity_variables(name, weight_hp)
 
   def __call__(self, inputs):
     w = self.sparsifiy(
@@ -98,7 +98,7 @@ def setUp(self):
           ),
       ),
   )
-  def test_create_aux_variables(self, mode_name, mode):
+  def test_create_sparsity_variables(self, mode_name, mode):
     sparsity_p = pax_fiddle.Config(
         SparsityHParams,
         sparsity_type=SparsityType.STRUCTURED_NM,