def testSwitchGradWithWeakTypeMismatch(self):
    dtype = jnp.ones(1).dtype
    dtype = (jnp.float32 if (dtype == jnp.float32) else jnp.float64)
    branches = [(lambda x: x), (lambda x: (x + dtype(1)))]

    def f_ref(x):
        i = x.astype(jnp.int32)
        return branches[i](x)

    def f(x):
        return lax.switch(x.astype(jnp.int32), branches, x)
    for x in [0.0, 1.0]:
        ans = api.grad(f)(x)
        expected = api.grad(f_ref)(x)
        self.assertAllClose(ans, expected, check_dtypes=False)
def testSwitchGradWithWeakTypeMismatch(self):
    dtype = jnp.ones(1).dtype
    dtype = (jnp.float32 if (dtype == jnp.float32) else jnp.float64)
    branches = [(lambda x: x), (lambda x: (x + dtype(1)))]

    def f_ref(x):
        i = x.astype(jnp.int32)
        return branches[i](x)

    def f(x):
        return lax.switch(x.astype(jnp.int32), branches, x)
    for x in [0.0, 1.0]:
        ans = api.grad(f)(x)
        expected = api.grad(f_ref)(x)
        self.assertAllClose(ans, expected, check_dtypes=False)