Remove underscored funcs

Signed-off-by: Jacob Platin <jacobplatin@google.com>

Author: jrplatin

tests/test_quantization.py

@@ -647,10 +647,10 @@ def setUp(self):
         self.calibration_method = 'max'
 
     @patch(
-        'tpu_commons.models.jax.utils.quantization.quantization_utils.ptq._create_quantized_param'
+        'tpu_commons.models.jax.utils.quantization.quantization_utils.ptq.create_quantized_param'
     )
     def test_manually_quantize_qwix_weight(self, mock_create_param):
-        """Test that manually_quantize_qwix_weight calls ptq._create_quantized_param correctly."""
+        """Test that manually_quantize_qwix_weight calls ptq.create_quantized_param correctly."""
         quantize_qwix.manually_quantize_qwix_weight(
             weight=self.weight,
             qtype=self.qtype,
@@ -672,12 +672,12 @@ def test_manually_quantize_qwix_weight(self, mock_create_param):
                          self.calibration_method)
 
     @patch(
-        'tpu_commons.models.jax.utils.quantization.quantization_utils.ptq._quantize_act'
+        'tpu_commons.models.jax.utils.quantization.quantization_utils.ptq.quantize_act'
     )
     @patch('qwix.pallas.get_current_rule')
     def test_manually_quantize_qwix_activation(self, mock_get_rule,
                                                mock_quantize_act):
-        """Test that manually_quantize_qwix_activation calls ptq._quantize_act correctly."""
+        """Test that manually_quantize_qwix_activation calls ptq.quantize_act correctly."""
         mock_rule = MagicMock()
         mock_rule.act_static_scale = False
         mock_get_rule.return_value = mock_rule

tpu_commons/models/jax/utils/quantization/quantization_utils.py

@@ -527,7 +527,7 @@ def manually_quantize_qwix_weight(weight: jax.Array, qtype: jnp.dtype,
         tiled_axes=tiled_axes,
         calibration_method=calibration_method)
 
-    return ptq._create_quantized_param(weight, how_to_quantize)
+    return ptq.create_quantized_param(weight, how_to_quantize)
 
 
 def manually_quantize_qwix_activation(inputs: jax.Array, rule_name: str,
@@ -561,7 +561,7 @@ def manually_quantize_qwix_activation(inputs: jax.Array, rule_name: str,
     # channelwise_axes should be set to (a subset of) non-contraction axes. e.g.
     # for ragged_dot [m, k] x [g, k, n], they are [0] and [0, 2]
     # TODO (jacobplatin): add support for `act_name`
-    return ptq._quantize_act(inputs, lhs_how, rule, "")
+    return ptq.quantize_act(inputs, lhs_how, rule, "")
 
 
 def get_quant_dtype_from_qwix_config(