@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_start_shape={}_stop_shape={}_num={}_endpoint={}_dtype={}'.format(start_shape, stop_shape, num, endpoint, dtype), 'start_shape': start_shape, 'stop_shape': stop_shape, 'num': num, 'endpoint': endpoint, 'dtype': dtype, 'rng_factory': rng_factory} for start_shape in [(), (2,), (2, 2)] for stop_shape in [(), (2,), (2, 2)] for num in [0, 1, 2, 5, 20] for endpoint in [True, False] for dtype in (inexact_dtypes + [None]) for rng_factory in [jtu.rand_default])))
def testGeomspace(self, start_shape, stop_shape, num, endpoint, dtype, rng_factory):
    rng = rng_factory()
    tol = {onp.float16: 0.004, onp.float32: 0.002, onp.complex128: 1e-14}

    def args_maker():
        'Test the set of inputs onp.geomspace is well-defined on.'
        (start, stop) = self._GetArgsMaker(rng, [start_shape, stop_shape], [dtype, dtype])()
        (start, stop) = jnp.broadcast_arrays(start, stop)
        if (dtype in complex_dtypes):
            return (start, stop)
        start = ((start * jnp.sign(start)) * jnp.sign(stop))
        return (start, stop)
    (start, stop) = args_maker()
    ndim = len(onp.shape((start + stop)))
    for axis in range((- ndim), ndim):

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

        def onp_op(start, stop):
            start = (start.astype(onp.float32) if (dtype == jnp.bfloat16) else start)
            stop = (stop.astype(onp.float32) if (dtype == jnp.bfloat16) else stop)
            return onp.geomspace(start, stop, num, endpoint=endpoint, dtype=(dtype if (dtype != jnp.bfloat16) else onp.float32), axis=axis).astype(dtype)
        self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=False, tol=tol)
        if (dtype in (inexact_dtypes + [None])):
            self._CompileAndCheck(jnp_op, args_maker, check_dtypes=False, atol=tol, rtol=tol)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_start_shape={}_stop_shape={}_num={}_endpoint={}_dtype={}'.format(start_shape, stop_shape, num, endpoint, dtype), 'start_shape': start_shape, 'stop_shape': stop_shape, 'num': num, 'endpoint': endpoint, 'dtype': dtype, 'rng_factory': rng_factory} for start_shape in [(), (2,), (2, 2)] for stop_shape in [(), (2,), (2, 2)] for num in [0, 1, 2, 5, 20] for endpoint in [True, False] for dtype in (inexact_dtypes + [None]) for rng_factory in [jtu.rand_default])))
def testGeomspace(self, start_shape, stop_shape, num, endpoint, dtype, rng_factory):
    rng = rng_factory()
    tol = {onp.float16: 0.004, onp.float32: 0.002, onp.complex128: 1e-14}

    def args_maker():
        'Test the set of inputs onp.geomspace is well-defined on.'
        (start, stop) = self._GetArgsMaker(rng, [start_shape, stop_shape], [dtype, dtype])()
        (start, stop) = jnp.broadcast_arrays(start, stop)
        if (dtype in complex_dtypes):
            return (start, stop)
        start = ((start * jnp.sign(start)) * jnp.sign(stop))
        return (start, stop)
    (start, stop) = args_maker()
    ndim = len(onp.shape((start + stop)))
    for axis in range((- ndim), ndim):

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

        def onp_op(start, stop):
            start = (start.astype(onp.float32) if (dtype == jnp.bfloat16) else start)
            stop = (stop.astype(onp.float32) if (dtype == jnp.bfloat16) else stop)
            return onp.geomspace(start, stop, num, endpoint=endpoint, dtype=(dtype if (dtype != jnp.bfloat16) else onp.float32), axis=axis).astype(dtype)
        self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=False, tol=tol)
        if (dtype in (inexact_dtypes + [None])):
            self._CompileAndCheck(jnp_op, args_maker, check_dtypes=False, atol=tol, rtol=tol)