Fix mixed precision compatibility issues in Bayesian Neural Networks example

[samudraneel05]

Fixes #1860

This PR fixes mixed precision/dtype compatibility issues in the Bayesian Neural Networks example - examples/keras_recipes/bayesian_neural_networks.py that caused the example to fail when mixed precision is enabled.

Problem

The example was failing with dtype-related errors when using TensorFlow Probability layers. The layers were not explicitly specifying dtypes, leading to type mismatches, errors and other dtype incompatibility issues.

Changes Made

• Explicitly set dtype=tf.float32 for all TensorFlow Probability distribution parameters:
  • prior(): Added dtype specification to tf.zeros() and tf.ones() in MultivariateNormalDiag
  • posterior(): Added dtype to VariableLayer and MultivariateNormalTriL
  • create_bnn_model(): Added dtype to DenseVariational layer
  • create_probablistic_bnn_model(): Added dtype to DenseVariational and Dense layers, and IndependentNormal output layer
  • negative_loglikelihood(): Added explicit casting of targets to tf.float32

Testing

Ran tests on my local machine (Macbook M2 with Python 3.13), converted .py to ipynb and ran the full notebook without errors on mixed precision data.

[gemini-code-assist]

Summary of Changes

Hello @samudraneel05, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses and resolves critical data type compatibility issues within the Bayesian Neural Networks example, particularly when running with mixed precision. By explicitly enforcing tf.float32 for TensorFlow Probability layers and operations, the example now executes without dtype-related errors, ensuring its robustness and correct functionality under various precision settings. This change enhances the example's reliability and broadens its applicability for users leveraging mixed precision training.

Highlights

• Mixed Precision Compatibility: Resolved dtype incompatibility issues in the Bayesian Neural Networks example, specifically when mixed precision is enabled, which previously caused the example to fail.
• Explicit Dtype Specification: Explicitly set dtype=tf.float32 across various TensorFlow Probability layers and operations, including tf.zeros(), tf.ones(), VariableLayer, MultivariateNormalTriL, DenseVariational, Dense, and IndependentNormal layers.
• Target Casting: Added explicit casting of targets to tf.float32 within the negative_loglikelihood function to ensure consistent data types.
• Keras Import Update: Updated Keras imports from tensorflow.keras to tf_keras for broader compatibility.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses mixed precision compatibility issues in the Bayesian Neural Networks example by explicitly setting dtype=tf.float32 in various TensorFlow Probability and Keras layers. The changes are generally correct and should resolve the reported failures when mixed precision is enabled. I've provided a few suggestions to improve the implementation by using the dtype parameter passed to functions instead of hardcoding tf.float32, which makes the code more robust and maintainable. I also have a question regarding the change of imports from tensorflow.keras to tf_keras, as it's not mentioned in the description and seems to deviate from the Keras 3 migration guide.

[samudraneel05]

Hey @sachinprasadhs, if the scope of the fix appears small, I'm happy to also update the .ipynb and the .md files.