def jnp_op(start, stop):
    return jnp.geomspace(start, stop, num, endpoint=endpoint, dtype=dtype, axis=axis)
def jnp_op(start, stop):
    return jnp.geomspace(start, stop, num, endpoint=endpoint, dtype=dtype, axis=axis)