Tags: google/qwix
Tags
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