Update chex.assert_type to check concrete types instead of just asserting that the type is a floating/integer sub-type.

[copybara-service]

Update chex.assert_type to check concrete types instead of just asserting that the type is a floating/integer sub-type.

Previously, assert_type would only check that the input was of the same parent type. For example:

x = np.ones((1,), dtype=np.float32)
chex.assert_type(x, np.float64)  # Succeeds
chex.assert_type(x, np.int32)  # Fails.

Instead, if a concrete dtype is provided we check that the input has the same type. If float or np.floating is provided, we continue to only assert that the input is the same parent.

x = np.ones((1,), dtype=np.float32)
chex.assert_type(x, np.float64) # Fails
chex.assert_type(x, float) # Succeeds.

[stephen-huan]

@tomwardio distrax/_src/utils/hmm_test.py also uses chex.assert_type here

    with self.subTest("states"):
      chex.assert_type(states, jnp.int32)
      chex.assert_shape(states, (length,))


    with self.subTest("observations"):
      chex.assert_type(obs, model.obs_dist.dtype)
      chex.assert_shape(obs, (length, *model.obs_dist.event_shape))

and here.

    with self.subTest("alphas"):
      chex.assert_type(alphas, jnp.float32)
      chex.assert_shape(alphas, (length, num_states))


    with self.subTest("betas"):
      chex.assert_type(betas, jnp.float32)
      chex.assert_shape(betas, (length, num_states))


    with self.subTest("marginals"):
      chex.assert_type(marginals, jnp.float32)
      chex.assert_shape(marginals, (length, num_states))


    with self.subTest("log_prob"):
      chex.assert_type(log_prob, jnp.float32)
      chex.assert_shape(log_prob, ())