import jax.numpy as jnp
from jax import lax
import optax
import chex


def _onehot(labels: chex.Array, num_classes: int) -> chex.Array:
    x = labels[..., None] == jnp.arange(num_classes).reshape((1,) * labels.ndim + (-1,))
    x = lax.select(x, jnp.ones(x.shape), jnp.zeros(x.shape))
    return x.astype(jnp.float32)


def p_contrastive_loss(ss: chex.Array, tt: chex.Array, axis: str = 'device') -> chex.Array:
    per_shard_targets = tt.shape[0]
    per_sample_targets = int(tt.shape[0] / ss.shape[0])
    labels = jnp.arange(0, per_shard_targets, per_sample_targets) + per_shard_targets * lax.axis_index(axis)

    tt = lax.all_gather(tt, axis).reshape((-1, ss.shape[-1]))
    scores = jnp.dot(ss, jnp.transpose(tt))

    return optax.softmax_cross_entropy(scores, _onehot(labels, scores.shape[-1]))

[(9, 'jax.numpy.ones', 'jnp.ones', True, 'import jax.numpy as jnp\n'), (9, 'jax.numpy.zeros', 'jnp.zeros', True, 'import jax.numpy as jnp\n'), (16, 'jax.numpy.arange', 'jnp.arange', True, 'import jax.numpy as jnp\n'), (19, 'jax.numpy.transpose', 'jnp.transpose', True, 'import jax.numpy as jnp\n'), (16, 'jax.lax.axis_index', 'lax.axis_index', False, 'from jax import lax\n'), (18, 'jax.lax.all_gather', 'lax.all_gather', False, 'from jax import lax\n'), (8, 'jax.numpy.arange', 'jnp.arange', True, 'import jax.numpy as jnp\n')]