[Bug] Missing float32 cast in TopP and Random samplers causes crashes in Mixed Precision

[AlejandroPG06]

Description

I have identified a potential stability issue when using TopPSampler and RandomSampler under a mixed precision policy (float16) with the TensorFlow backend.

In src/samplers/top_k_sampler.py, the inputs to random.categorical are explicitly cast to "float32" to prevent runtime errors, as noted in the source code comments:

tf does not support half precision multinomial sampling, so make sure we have full precision here.

However, this safeguard is missing in:

1. src/samplers/top_p_sampler.py
2. src/samplers/random_sampler.py

Impact

Running generation tasks with these samplers in a mixed-precision environment may lead to runtime exceptions or numerical instability on backends that do not support float16 for multinomial sampling operations.

Proposed Fix

I have applied the same ops.cast(..., "float32") logic used in TopKSampler to both TopPSampler and RandomSampler to ensure consistency and stability across all sampling strategies.

[sachinprasadhs]

Could you please provide sample reproducible code to replicate the issue reported?
I could not find an issue while trying to reproduce this.

[AlejandroPG06]

Thank you for taking the time to review this issue! I understand the test in the Colab runs successfully, but I'd like to clarify why this doesn't invalidate the underlying consistency issue.

The Core Issue: Code Inconsistency, Not Just Runtime Errors

This issue is fundamentally about code consistency and best practices rather than just runtime failures. Let me demonstrate with evidence from KerasHub's own source code.

Evidence from TopKSampler Source Code

In src/samplers/top_k_sampler.py, the developers explicitly documented the float16 issue:

def __call__(self, prompt, next_token_logits):
    # ... top-k filtering ...
    
    # tf does not support half precision multinomial sampling, so make
    # sure we have full precision here.
    next_token_logits = ops.cast(next_token_logits, "float32")
    
    next_token = ops.random.categorical(
        next_token_logits,
        1,
        seed=self.seed,
    )
    return ops.squeeze(next_token, axis=-1)

Key observation: The comment clearly states: "tf does not support half precision multinomial sampling"

This is a documented limitation that the TopKSampler developers were aware of and explicitly protected against.

The Inconsistency

However, when we examine RandomSampler and TopPSampler:

RandomSampler (src/samplers/random_sampler.py):

def __call__(self, prompt, next_token_logits):
    # NO float32 casting
    next_token = ops.random.categorical(
        next_token_logits,  # ← Could be float16!
        1,
        seed=self.seed,
    )
    return ops.squeeze(next_token, axis=-1)

TopPSampler (src/samplers/top_p_sampler.py):

def __call__(self, prompt, next_token_logits):
    # ... top-p filtering ...
    # NO float32 casting
    next_token = ops.random.categorical(
        next_token_logits,  # ← Could be float16!
        1,
        seed=self.seed,
    )
    return ops.squeeze(next_token, axis=-1)

Both are missing the ops.cast(next_token_logits, "float32") safeguard that TopKSampler has.

Why Your Test May Still Pass

The test in your Colab may pass for several reasons:

1. TensorFlow Auto-casting: Recent TF versions may implicitly cast to float32 in some scenarios
2. Model Architecture: GPT2's specific architecture might output float32 logits despite mixed precision
3. Version-Specific Behavior: Different TF/Keras versions handle dtype differently
4. Hardware-Specific: Behavior can vary between CPU/GPU and different GPU architectures

However, relying on implicit behavior is fragile and goes against the documented best practice already established in TopKSampler.

Why This Should Be Fixed

Even if it doesn't always cause runtime errors, this inconsistency should be addressed because:

1. Code Consistency: All samplers should follow the same best practices
2. Future-Proofing: Protects against TensorFlow version changes
3. Documentation: TopKSampler's comment explicitly warns about this issue
4. Maintainability: Reduces cognitive load - all samplers work the same way
5. Defensive Programming: Explicit is better than implicit

Proposed Solution

Apply the same safeguard from TopKSampler to the other samplers:

# In RandomSampler and TopPSampler __call__ methods, before categorical sampling:
next_token_logits = ops.cast(next_token_logits, "float32")

This is a minimal, safe change that:

• Aligns with TopKSampler's documented pattern
• Ensures consistent behavior across all samplers
• Protects against future TF changes
• Follows explicit-over-implicit principle

Conclusion

I believe the issue here is about maintaining code consistency with a pattern that KerasHub developers themselves documented as necessary in TopKSampler.

If TopKSampler needs ops.cast(..., "float32") to prevent float16 issues with tf.random.categorical, then RandomSampler and TopPSampler should have the same protection for consistency and robustness.

Would you be open to this change for consistency?

---

Additional Resources:

• I've prepared some detailed scripts demonstrating the source code inconsistency
• I can provide additional evidence if needed

Thank you for considering this! 🙏

[sachinprasadhs]

The issue is already handled in the base class here https://github.com/keras-team/keras-hub/pull/1242/changes.

The changes in the PR which you have suggested becomes redundant and extra guardrail which is not required, since the issue is not present with the mixed precision float16.