0.1.6

What's Changed

• update pyproject version to match 0.1.5 by @jshin1394 in #188
• Add JAX op links to the Qwix core API page via docstring. by @jshin1394 in #189
• Register flax.nnx.BatchNorm.__call__ in odml_ops.py to support nnx.BatchNorm fusion. by @copybara-service[bot] in #190
• [QwixConvGeneral] Add out_sharding parameter to Qwix conv_general_dilated functions. by @copybara-service[bot] in #193
• Add option to use reciprocal for quantization instead of division. by @copybara-service[bot] in #196
• [QwixFlaxUtil] Replace nn.unbox with local implementation of unbox in find_param in flax_util.py. by @copybara-service[bot] in #197

• Fix GPTQ producing invalid QArray when subchannel quantization not specified by @copybara-service[bot] in #198

• [QwixOpId] Improve op_id generation for Linen and NNX modules. by @copybara-service[bot] in #200
• [QwixGptqEinsum] Add einsum to GptqCalibrationProvider and support attention layer weight GPTQ. by @copybara-service[bot] in #201
• Adds AWQ (Activation-aware Weight Quantization) support. by @copybara-service[bot] in #205
• properly set LoRA alpha by @copybara-service[bot] in #206
• [QwixCitation] Add citation information to the Qwix README. by @copybara-service[bot] in #207
• Refactor interception mechanism using a singleton _InterceptionManager. by @copybara-service[bot] in #208
• Fix argument passing in intercepted function calls. by @copybara-service[bot] in #210
• Adds dev setup documentation by @copybara-service[bot] in #212
• Rename sharding_names to out_sharding in NNX Variable metadata by @copybara-service[bot] in #199
• [QwixEinsum] Refactor Einsum to handle symbolic input shapes. by @copybara-service[bot] in #215
• [QwixNnxDataType] Register Qwix metadata and array structures as NNX data types. by @copybara-service[bot] in #218
• [QwixStructuralInterception] Enable structural binding for PrimitiveBindOp to support transpose op removal and allow for robust ODML metadata propagation. by @copybara-service[bot] in #211
• averaging_test, gitignore agents by @copybara-service[bot] in #219
• [QwixOdmlTagPropagation] Fix ODML tag propagation by adding gather and split to _VALUE_PRESERVING_PRIMITIVES. by @copybara-service[bot] in #220
• Change logging level in qconfig.py from INFO to DEBUG. by @copybara-service[bot] in #221
• Simplify uniform random numbers. Remove unnecessary tests. by @copybara-service[bot] in #222
• change defaults for Dropout and BatchNorm by @copybara-service[bot] in #224
• [QwixOdmlOp] Update odml_op to 1) handle every leaf of pytree output activations for fake quantization and 2) allow ops to veto quantization requests from previous ops if the op is excluded from quantization. by @copybara-service[bot] in #223
• [QwixNumerics] Improve docstring for qwix numerics functions. by @copybara-service[bot] in #226
• Fix channelwise contracting axes in _fast_dot_general. by @copybara-service[bot] in #228
• Adds sparsity functionality for QWIX. by @copybara-service[bot] in #225
• [QwixMxfpNumerics] Verify XLA handling of MXFP data types by testing jax.nn.scaled_matmul across different platforms. by @copybara-service[bot] in #227
• Add process_prequantized_params to Qwix PTQ provider. by @copybara-service[bot] in #231
• Add integration test for Orbax + Qwix Safetensors loading by @copybara-service[bot] in #233
• [Qwix2dBlockScale] Enable native loading of 2D block-scaled checkpoints by @copybara-service[bot] in #235
• Add rename transform support to load_nested_safetensors in Orbax + Qwix integration by @copybara-service[bot] in #239
• Remove sparsity_rule from HowToQuantize by @BirdsOfAFthr in #241
• [QwixLogging] Standardize logging to absl in Qwix tests by @copybara-service[bot] in #245
• [QwixInterception] Selective disable_jit for XLA and ODML providers. by @copybara-service[bot] in #230
• pallas_call isn't being intercepted correctly. by @copybara-service[bot] in #246
• Add init file to contrib folder. by @copybara-service[bot] in #248
• [QwixInterception] Selective disable_jit for XLA and ODML providers, improved with stable Interceptor ID. by @copybara-service[bot] in #247
• Add tracking for unused quantization rules. by @copybara-service[bot] in #244
• [QwixPreQuantizedCheckpoint] Improve qwix.process_prequantized_params to handle checkpoint source of truth conversions. by @copybara-service[bot] in #250
• [JAX] Fix code that compares PartitionSpecs with tuples. by @copybara-service[bot] in #251
• Refactor residual handling in Qwix. by @copybara-service[bot] in #243
• [QwixPreQuantizedCheckpoint] Simplify process_prequantized_params error handling. by @copybara-service[bot] in #252
• This CL, while addressing a crash in _fast_dot_general when using per-element scales (tile_size=1) on contracting axes, fundamentally changed the behavior of the core utility qarray.get_tiled_axes. Previously, get_tiled_axes was understood to report only subchannel tiling (tile_size > 1). The change to include per-element scaled axes (tile_size=1) for all axis types, including contracting ones, introduces complexity and potential unintended side effects across the Qwix codebase, as other modules might implicitly rely on the old behavior. by @copybara-service[bot] in #253
• Consolidated qwix improvements for LoRA and sharding support. by @copybara-service[bot] in #242
• quantized matmul kernel by @copybara-service[bot] in #240

New Contributors

• @BirdsOfAFthr made their first contribution in #241

Full Changelog: 0.1.5...0.1.6

0.1.5

What's Changed

• Allow Qwix logo to be transparent. by @jshin1394 in #171
• Reland: Fix the RNG issue for LoRA by @copybara-service[bot] in #172
• [QwixQuantizedTraining] Clip Gradients for Out-of-Bound Activations and Weights in dot_general_qt. by @copybara-service[bot] in #173
• update readthedocs contents to include latest changes by @jshin1394 in #174
• Explicitly look for ptq.WithAux in find_param by @copybara-service[bot] in #181
• Allow scale dtype to be float16. by @copybara-service[bot] in #178
• Relax einsum tests. by @copybara-service[bot] in #183
• Generalize GPTQ to support multi-dimensional weights. by @copybara-service[bot] in #184
• [QwixPallas] Add more detail to dot_general.loop_dot_general docstring, which is used inside qwix pallas.pallas_call. by @copybara-service[bot] in #185
• Add API documentation to Qwix readthedocs by @jshin1394 in #186

Full Changelog: 0.1.4...0.1.5