keras/src/backend/jax/numpy.py/transpose
def transpose(x, axes=None):
    x = convert_to_tensor(x)
    if isinstance(x, jax_sparse.BCOO):
        num_dims = len(x.shape)
        if axes is None:
            permutation = tuple(range(num_dims)[::-1])
        else:
            permutation = []
            for a in axes:
                a = canonicalize_axis(a, num_dims)
                permutation.append(a)
        return jax_sparse.bcoo_transpose(x, permutation=permutation)
    return jnp.transpose(x, axes=axes)