Conversation
There was a problem hiding this comment.
Code Review
This pull request fixes a bug in the JAX backend's convert_to_tensor where floatx would override an input's dtype when floatx is set to "bfloat16". It ensures that the bfloat16 fast path is only triggered when dtype is explicitly requested, and adds comprehensive regression tests to verify this behavior. The review feedback suggests simplifying the condition in convert_to_tensor by removing the redundant standardize_dtype call, as dtype is already standardized earlier in the function.
41026a5 to
b1cf8ac
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #23683 +/- ##
==========================================
- Coverage 86.96% 84.40% -2.57%
==========================================
Files 484 484
Lines 71754 71754
Branches 11840 11840
==========================================
- Hits 62400 60562 -1838
- Misses 6224 8185 +1961
+ Partials 3130 3007 -123
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This PR fixes bfloat16 fast path in the JAX backend was guarded by
standardize_dtype(dtype) == "bfloat16". Sincestandardize_dtype(None)returnsfloatx(), that branch fired for every call that omitteddtypeonce
floatxwas "bfloat16", and the followingastype(None)resolved to float32.This change ensures the branch only runs on an explicit bfloat16 request by checking
dtype is not None.Fixes: #23679
Contributor Agreement
Please review our PR Contribution Policy and AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:
Fixes #xxxabove).Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed. PRs without a linked, assigned issue will be converted to draft.