JAX, kerangka komputasi berbasis array dari Google, telah menjadi pilihan utama bagi peneliti dan praktisi machine learning karena kemampuannya dalam mempercepat perhitungan pada CPU, GPU, dan TPU. Secara bawaan, JAX menganggap array sebagai tipe primitif, sehingga fungsi yang ditransformasi selalu menerima dan menghasilkan array. Untuk menangani struktur data yang kompleks, JAX menyediakan mekanisme pytree yang secara transparan meratakan kontainer menjadi leaf‑leaf array. Namun, transparansi ini seringkali membatasi kemampuan untuk memodelkan data dengan invariants tertentu, seperti array kuantisasi yang terdiri dari nilai diskrit dan skala kontinyu yang terkait erat.
Ketika bekerja dengan array kuantisasi, misalnya nilai int8 yang disertai skala float per baris, pendekatan pytree kehilangan semantik penting. Kedua komponen tersebut harus selalu konsisten: skala harus memiliki bentuk yang sesuai dengan nilai kuantisasi, dan keduanya tidak boleh dipisahkan saat transformasi. Selain itu, tipe tangent dari pytree dipaksa mengikuti tipe leaf, sehingga tidak ada cara untuk mendefinisikan turunan yang bermakna bagi array kuantisasi. Hal inilah yang mendorong munculnya hijax, sebuah ekstensi eksperimental JAX yang memungkinkan pembuatan tipe data khusus dengan identitas sendiri.
Hijax types (atau “tipe hi”) memungkinkan pengembang untuk membuat tipe baru dengan subclassing HiType. Implementasi tipe ini mencakup metode lo_ty yang mendeskripsikan tipe array lojax yang menjadi dasar, serta lower_val dan raise_val untuk konversi nilai. Setelah tipe didefinisikan, kelas Python yang membawa nilai tipe tersebut harus didaftarkan menggunakan register_hitype. Dengan demikian, tipe baru dapat muncul sebagai satu entitas utuh dalam jaxprs, mendukung invariants, tangent khusus, batching melalui vmap, dan sharding eksplisit.
Contoh konkret dapat dilihat pada tipe array kuantisasi (QArray). Implementasi dimulai dengan kelas dataclass QArray yang menyimpan qvalue (int8) dan scale (float32). Kelas tipe QArrayTy mewarisi HiType, menyimpan shape dan sharding dari NamedSharding. Metode lo_ty kemudian mengembalikan dua ShapedArray: satu untuk qvalue dengan sharding yang diberikan, dan satu untuk scale dengan sharding yang berasal dari sharding qvalue dengan menghapus sumbu terakhir. lower_val dan raise_val bertugas mengonversi instance QArray ke daftar array dan sebaliknya.
Untuk mendukung diferensiasi otomatis, tipe hi harus menyediakan to_tangent_aval yang mendefinisikan tipe tangent yang sesuai, misalnya float32 kontinu yang mendekati nilai kuantisasi. Selain itu, aturan VJP dan JVP harus diimplementasikan pada primitif hijax yang menggunakan tipe baru ini, sehingga transformasi seperti grad dapat berfungsi tanpa kehilangan semantik kuantisasi.
Dukungan vmap memerlukan implementasi dec_rank dan inc_rank pada tipe hi, bersama dengan MappingSpec khusus yang mendefinisikan cara batching dilakukan. Argumen tipe hi yang dimetakan memerlukan axis_size dan in_axes/out_axes yang eksplisit, sehingga transformasi pemetaan dapat diterapkan dengan benar tanpa mengorbankan invariants internal.
Integrasi sharding dalam tipe hi dilakukan dengan menyimpan informasi partisi pada tipe, misalnya field NamedSharding. Field ini kemudian digunakan dalam lo_ty untuk memberi label pada array komponen dengan sharding yang tepat, sehingga tipe tersebut dapat berpartisipasi dalam mode sharding eksplisit JAX. Pendekatan ini memungkinkan array kuantisasi untuk didistribusikan secara merata di seluruh perangkat, yang sangat berguna untuk pelatihan model besar di lingkungan multi-GPU.
Meskipun hijax masih bersifat eksperimental dan API-nya dapat berubah, konsep ini membuka peluang baru bagi pengembang untuk memodelkan data dengan semantik khusus di dalam JAX. Di Indonesia, di mana sumber daya komputasi semakin melimpah dan komunitas AI/ML terus berkembang, kemampuan untuk mendefinisikan tipe data khusus dapat mempercepat penelitian dalam bidang kuantisasi, kompresi, dan pemrosesan data sensitif privasi. Dengan memanfaatkan hijax, peneliti lokal dapat membangun pipeline yang lebih efisien dan akurat, sekaligus berkontribusi pada ekosistem JAX secara global.