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])()
      # onp.geomspace can't handle differently ranked tensors
      # w. negative numbers!
      start, stop = lnp.broadcast_arrays(start, stop)
      if dtype in complex_dtypes:
        return start, stop
      # to avoid NaNs, non-complex start and stop cannot
      # differ in sign, elementwise
      start = start * lnp.sign(start) * lnp.sign(stop)
      return start, stop
Test the set of inputs onp.geomspace is well-defined on.
trax/tf_numpy/jax_tests/lax_numpy_test.py