@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix('', shapes, (((onp.bool_,) * n) + dtypes)), 'rng_factory': jtu.rand_default, 'shapes': shapes, 'dtypes': dtypes} for n in range(0, 3) for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(all_shapes, ((2 * n) + 1))) for dtypes in CombosWithReplacement(all_dtypes, (n + 1)))))
def testSelect(self, rng_factory, shapes, dtypes):
    rng = rng_factory()
    n = (len(dtypes) - 1)

    def args_maker():
        condlist = [rng(shape, onp.bool_) for shape in shapes[:n]]
        choicelist = [rng(shape, dtype) for (shape, dtype) in zip(shapes[n:(- 1)], dtypes[:n])]
        default = rng(shapes[(- 1)], dtypes[(- 1)])
        return (condlist, choicelist, default)

    def onp_fun(condlist, choicelist, default):
        choicelist = [(x if (jnp.result_type(x) != jnp.bfloat16) else x.astype(onp.float32)) for x in choicelist]
        dtype = jnp.result_type(default, *choicelist)
        return onp.select(condlist, [onp.asarray(x, dtype=dtype) for x in choicelist], onp.asarray(default, dtype=dtype))
    self._CheckAgainstNumpy(onp_fun, jnp.select, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp.select, args_maker, check_dtypes=True, rtol={onp.float64: 1e-07, onp.complex128: 1e-07})
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix(, shapes, (((onp.bool_,) * n) + dtypes)), 'rng_factory': jtu.rand_default, 'shapes': shapes, 'dtypes': dtypes} for n in range(0, 3) for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(all_shapes, ((2 * n) + 1))) for dtypes in CombosWithReplacement(all_dtypes, (n + 1)))))
def testSelect(self, rng_factory, shapes, dtypes):
    rng = rng_factory()
    n = (len(dtypes) - 1)

    def args_maker():
        condlist = [rng(shape, onp.bool_) for shape in shapes[:n]]
        choicelist = [rng(shape, dtype) for (shape, dtype) in zip(shapes[n:(- 1)], dtypes[:n])]
        default = rng(shapes[(- 1)], dtypes[(- 1)])
        return (condlist, choicelist, default)

    def onp_fun(condlist, choicelist, default):
        choicelist = [(x if (jnp.result_type(x) != jnp.bfloat16) else x.astype(onp.float32)) for x in choicelist]
        dtype = jnp.result_type(default, *choicelist)
        return onp.select(condlist, [onp.asarray(x, dtype=dtype) for x in choicelist], onp.asarray(default, dtype=dtype))
    self._CheckAgainstNumpy(onp_fun, jnp.select, args_maker, check_dtypes=False)
    self._CompileAndCheck(jnp.select, args_maker, check_dtypes=True, rtol={onp.float64: 1e-07, onp.complex128: 1e-07})