Tags: google/qwix
Tags
Add Hadamard rotation utilities to qwix/contrib. This change introduces a `hadamard_rot` module containing functions to generate Hadamard matrices with optional random sign flips and apply them to tensors. The core functionality is `hadamard_rotate_inputs`, which transforms both an activation and a weight matrix such that their product remains invariant. This is useful for techniques like Hadamard-based quantization. PiperOrigin-RevId: 933209914
Switch to an alternative way to discover weight_name This change introduces `flax_util.find_param` to dynamically find the name of a parameter within the current Flax/NNX module. This eliminates the need to store and propagate `weight_name` using `aux_data`, simplifying the interception logic in PTQ and LoRA providers. The `WithAux` class no longer stores `weight_name`. The possible problem with this approach is we need some heuristics to find the weight name, and could possibly make mistake when the activation has the same shape as the weight, or there are multiple params with the same shape. In reality this should be rare and we have added proper checks for them. PiperOrigin-RevId: 814293120