`get_current_rule` is incorrectly a noop

[SamuelMarks]

Expected Behavior

It should be an op, or at least a raise NotImplementedError

Actual Behavior

#75 by @liudangyi created this noop 5-days ago in 354ceb9:

def get_current_rule(op_name: str) -> QuantizationRule | None:
  """Returns the current quantization rule if intercepted, or None otherwise."""
  del op_name
  return None

https://github.com/google/qwix/blob/8f1a822/qwix/_src/qconfig.py#L92-L95

FWIW: process_model_output seems to be a similar function that is a noop stub.

Steps to Reproduce the Problem

Specifications

• Version: still true at latest main (354ceb9)

[jshin1394]

Hi Samual! Thank you for bringing this to our attention! It seems like this intercept function should override the implementation such that get_current_rule should return the quantizationRule correctly(https://github.com/google/qwix/blob/main/qwix/_src/qconfig.py#L128-L129).

Please let me know if the function isn't working as intended. l)

  def get_intercept_map(self) -> dict[str, Callable[..., Any]]:
    """Returns the intercept map for interception.wrap_func_intercepted."""
    # Common functions that are intercepted by all quantization providers.
    return {
        'qwix._src.qconfig.get_current_rule': (
            lambda op: self._get_current_rule_and_op_id(op, only_rule=True)[0]
        ),
        # Disable interception for ops in pallas_call.
        'jax.experimental.pallas.pallas_call': (
            lambda *args, **kwargs: interception.disable_interceptions(
                pl.pallas_call(*args, **kwargs)
            )
        ),
    }

[SamuelMarks]

Thanks @jshin1394 but what if the function is called explicitly like here https://github.com/AI-Hypercomputer/maxtext/pull/2065/files#diff-76040df4a764dcc58034ee6d01cff34ed87a2a29a7b983948d101691e46f5d26R696

If the function isn't being replaced it will always return None. Can an error instead be raised, because that stub function body should never be executed?

[liudangyi]

This is intended and the design of Qwix. Returning None here allows the get_current_rule to be called regardless of whether quantization is enabled or not. For maxtext, the use_qwix_quantization will soon default to True and get_current_rule won't be in any condition.

[SamuelMarks]

@liudangyi The other way of doing it is to rewrite it to use a function called identity (or similar name that indicates None is always returned).

Here is the error that told me this function was wrong:

MaxText/layers/moe.py:701:8: E1128: Assigning result of a function call, where the function returns None (assignment-from-none

Indicating the first line here:

        quantization_rule = qpl.get_current_rule("dot_general")
        if quantization_rule is not None:
          lhs_quantize_dtype = quantization_rule.act_qtype
          rhs_quantize_dtype = quantization_rule.weight_qtype

🤷