def METHOD_NAME(op, node, **kwargs):
    """Register a JAX implementation for `ARange`.
    `jax.numpy.arange` requires concrete values for its arguments. Here we check
    that the arguments are constant, and raise otherwise.
    TODO: Handle other situations in which values are concrete (shape of an array).
    """
    arange_args = node.inputs
    constant_args = []
    for arg in arange_args:
        if arg.owner and isinstance(arg.owner.op, Shape_i):
            constant_args.append(None)
        elif isinstance(arg, Constant):
            constant_args.append(arg.value)
        else:
            # TODO: This might be failing without need (e.g., if arg = shape(x)[-1] + 1)!
            raise NotImplementedError(ARANGE_CONCRETE_VALUE_ERROR)
    constant_start, constant_stop, constant_step = constant_args
    def arange(start, stop, step):
        start = start if constant_start is None else constant_start
        stop = stop if constant_stop is None else constant_stop
        step = step if constant_step is None else constant_step
        return jnp.arange(start, stop, step, dtype=op.dtype)
    return arange