def _maybe_assert_valid_sample(self, counts):
    """Check counts for proper shape, values, then return tensor version."""
    if not self.validate_args:
      return counts
    counts = distribution_util.embed_check_nonnegative_integer_form(counts)
    return distribution_util.with_dependencies([
        assert_util.assert_less_equal(
            counts,
            self.total_count,
            message="counts are not less than or equal to n."),
    ], counts)
Check counts for proper shape, values, then return tensor version.