北京
你正嘗試把模型權重量化到int8，好讓推理更快。一個量化數組本質上是一對數組：整數數組加上每行一個的浮點縮放系數。可當你把這個復合體塞進JAX的變換流水線，就會發現它被自動拆成兩片葉子，所有內部約束都消失，自動微分也抓瞎。
透明或許令人省心，但在這種場景下反而成了累贅。JAX標準pytree機制默認將所有容器扁平化，可有些數據天生應該保持一個整體，有自己的類型身份、自己的不變量，用戶也只能通過一組固定操作去生產和消費它，而不是隨意拆解或拼接。這正是hijax模塊瞄準的缺口。
hijax讓你可以基于HiType定義新的JAX類型，并指定它到底層普通數組（所謂的lojax）的轉換方式。你注冊一個Python值類，再編寫專屬原語，規定輸入輸出類型涉及這個新類型。如此一來，這個類型在jaxpr中只會呈現為一個原子值，不再是散開的數組葉子，其內部結構對外部變換完全隱藏。
以量化數組為例，你先子類化HiType，實現lo_ty來描述底層應當解成什么數組，然后提供lower_val和raise_val來完成值與底層數組的互相轉換。通過調用register_hitype把值類和類型綁定，接著編寫VJPHiPrimitive子類，在in_avals與out_aval中引用新類型。這些原語就是生產或消費該類型值的唯一手段，內建的不變性從此得到保證。
自動微分方面，你可以在類型上實現to_tangent_aval，讓切線類型與原始結構分道揚鑣。量化整數的梯度可能并不需要一對對應梯度，而需要一個獨立的浮點表示。在VJP或JVP規則里按此定義梯度傳播，恰好匹配量化操作的實際需求。
碰上vmap批處理時，自定義類型得自己說明如何增減一個批次維度。你實現dec_rank和inc_rank，再設計一個MappingSpec子類來刻畫映射規格，接著在批量規則里把它們串起來。帶軸大小和規格的顯式參數會告訴vmap怎么處理這類輸入輸出，批量化也就順理成章。
如果還要利用JAX顯式分片模式，干脆把分片信息（比如NamedSharding）直接嵌在類型字段中。lo_ty里消費這把信息，原語的類型推導規則里再傳播它，讓多設備并行也能認得你這個新類型。雖然整套hijax API還貼著實驗標簽，但它已經展現了一條清晰路徑，讓JAX從單純數組世界邁向可擴展的富類型生態。
特別聲明：以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布，本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.
