import awkward as ak
import numpy as np
import jax

ak.jax.register_and_check()

test_regulararray = ak.Array(
  [[1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0]],
  backend="jax"
)

test_regulararray_tangent = ak.Array(
  [[1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0]],
  backend="jax"
)

value_jvp, jvp_grad = jax.jvp(
  lambda x: ak.sort(x, axis=0), (test_regulararray,), (test_regulararray_tangent,)
)
# >> ValueError: Encountered Traced<float64[9]>with<JVPTrace> with
#   primal = Array([1., 4., 7., 2., 5., 8., 3., 6., 9.], dtype=float64)
#   tangent = Array([1., 4., 7., 2., 5., 8., 3., 6., 9.], dtype=float64) as an (invalid) input to the 'awkward_sort' Awkward C++ kernel. This kernel is not differentiable by the JAX backend.
# 
# This error occurred while calling
#
#    ak.sort(
#        <Array [[...], [...], [...]] type='3 * var * float64'>
#        axis = 0
#    )