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, 22, 9, 39), 'jax.numpy.ones', 'jnp.ones', ({(9, 31, 9, 38): 'x.shape'}, {}), '(x.shape)', True, 'import jax.numpy as jnp\n'), ((9, 41, 9, 59), 'jax.numpy.zeros', 'jnp.zeros', ({(9, 51, 9, 58): 'x.shape'}, {}), '(x.shape)', True, 'import jax.numpy as jnp\n'), ((16, 13, 16, 65), 'jax.numpy.arange', 'jnp.arange', ({(16, 24, 16, 25): '(0)', (16, 27, 16, 44): 'per_shard_targets', (16, 46, 16, 64): 'per_sample_targets'}, {}), '(0, per_shard_targets, per_sample_targets)', True, 'import jax.numpy as jnp\n'), ((19, 25, 19, 42), 'jax.numpy.transpose', 'jnp.transpose', ({(19, 39, 19, 41): 'tt'}, {}), '(tt)', True, 'import jax.numpy as jnp\n'), ((16, 88, 16, 108), 'jax.lax.axis_index', 'lax.axis_index', ({(16, 103, 16, 107): 'axis'}, {}), '(axis)', False, 'from jax import lax\n'), ((18, 9, 18, 33), 'jax.lax.all_gather', 'lax.all_gather', ({(18, 24, 18, 26): 'tt', (18, 28, 18, 32): 'axis'}, {}), '(tt, axis)', False, 'from jax import lax\n'), ((8, 29, 8, 52), 'jax.numpy.arange', 'jnp.arange', ({(8, 40, 8, 51): 'num_classes'}, {}), '(num_classes)', True, 'import jax.numpy as jnp\n')]