JAX membuat kemajuan pesat di bidang sains, terutama dalam rekayasa protein berbasis AI. Startup Escalante menggunakan JAX untuk memprediksi dampak obat terhadap tingkat ekspresi protein intraseluler. Para pendiri, April Schleck dan Nick Boyd, menggunakan transformasi fungsi composable JAX untuk memecahkan masalah ilmiah dalam rekayasa protein yang memerlukan optimasi multi-objektif.

Escalante bertujuan membangun model pembelajaran mesin yang mampu merancang obat dari awal. Namun, kurangnya dataset biologis yang spesifik menimbulkan tantangan tersendiri. Untuk mengatasi tantangan ini, mereka mengembangkan uji laboratorium untuk menghasilkan data yang diperlukan. Sifat fungsional JAX sangat penting untuk mengintegrasikan berbagai model AI ke dalam satu putaran optimasi, sebuah tugas yang lebih kompleks dalam kerangka kerja lain.

April Schreck menyoroti peran JAX dalam proses desain protein, dengan mengintegrasikan beberapa model ke dalam satu fungsi kerugian. Pendekatan ini memungkinkan optimasi urutan protein menggunakan jaringan saraf pra-terlatih yang serupa dengan DeepDream milik Google. Metode Escalante mengembangkan urutan acak menjadi protein yang dioptimalkan melalui siklus evaluasi model dan pembaruan urutan.

Nick Boyd menyoroti efisiensi JAX dalam mengompilasi dan membedakan kode kompleks, yang sangat penting bagi penelitian mereka. Penggunaan perangkat keras TPU dan kompiler XLA memungkinkan penskalaan tugas komputasi. Infrastruktur Escalante memanfaatkan Google Kubernetes Engine untuk menjalankan ribuan Spot TPU untuk pekerjaan optimasi, sebuah strategi hemat biaya yang mendukung tujuan ambisius mereka untuk merancang pengikat protein yang menangani seluruh proteom manusia.

Ekosistem JAX, termasuk pustaka seperti Equinox dan Optax, telah menjadi bagian penting penelitian Escalante, dan pustaka Mosaic , yang tersedia di GitHub , adalah bukti bahwa JAX memperluas pengaruhnya di luar GitHub Learning, yang mendorong era baru penemuan ilmiah dan pemrograman yang dapat dibedakan.