JAX on Christian Hotz-Behofsits

Multi-hot encoding in JAX

Fri, 19 Aug 2022 00:00:00 +0000

JAX offers a range of practical functions, also for data preparation. One of these is jax.nn.one_hot, which performs a classic one-hot encoding. Unfortunately, I was not able to find a suitable multi-hot equivalent for multi-label applications. However, it is also quite easy to implement the functionality directly in JAX:

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnames=("num_classes"))
def multi_hot(labels, num_classes: int):
 return jnp.take(jnp.eye(num_classes), jnp.array(labels), axis=0).sum(axis=0)

Multinomial Logistic Regression in JAX

Fri, 29 Jul 2022 00:00:00 +0000

Classifications are a classic machine learning problem we can tackle using logistic regression. If we distinguish between more than two classes, we call it a multinomial logistic regression. In this post, I will show how this can be done using JAX based on the well-known Fisher’s Iris dataset (every R user should be familiar with this one).

First, we have to load the required libraries and load the data. Since this is a classification, we have a set of predictors (aka. features) and a label for each sample.