Open
Conversation
Owner
CodersAcademy006
commented
Dec 2, 2025
CodersAcademy006
commented
- Initial plan
- Add fallback mechanism for large tensors in LaunchConvOpImpl
- Address code review feedback: add bounds checking and validation
- Address remaining code review feedback: improve recursion safety comments and edge case handling
- Fix XLA JIT compilation with Keras initializers and dynamic shapes (XLA JIT Compilation Fails with Keras Initializers and Dynamic Shapes tensorflow/tensorflow#105334)
- Add comprehensive documentation for XLA initializers fix
- Fix XLA JIT compilation with mixed-type dictionary keys (XLA JIT Compilation Fails with Mixed-Type Dictionary Keys tensorflow/tensorflow#105333)
- Delete FIX_SUMMARY_105334.md
- Delete tensorflow/python/ops/demo_xla_initializers_fix.py
- Rename test file for Keras initializers
- Delete tensorflow/python/util/demo_mixed_dict_keys.py
- Rename test_mixed_dict_keys.py to mixed_dict_keys_test.py
- Remove unrelated cuDNN batch-splitting fallback from conv_ops_impl.h
- Revert conv_ops_impl.h changes (remove unrelated cuDNN fallback)
- Fix _compute_fans: robust XLA-safe _to_int conversion and correct receptive field calculation
- Move keras_initializers_dynamic_shapes_test to tensorflow/python/keras/initializers
- Fix imports in mixed_dict_keys_test to use internal TensorFlow APIs
- Clean PR: remove accidental keras initializer test from mixed-dict-keys branch
Co-authored-by: CodersAcademy006 <104912634+CodersAcademy006@users.noreply.github.com>
Co-authored-by: CodersAcademy006 <104912634+CodersAcademy006@users.noreply.github.com>
…ents and edge case handling Co-authored-by: CodersAcademy006 <104912634+CodersAcademy006@users.noreply.github.com>
…ensorflow#105334) This commit fixes issue tensorflow#105334 where @tf.function(jit_compile=True) fails when using Keras initializers with dynamic shapes containing symbolic tensors. Root Cause: ----------- When XLA JIT compilation is enabled, tf.shape() returns symbolic tensors rather than concrete values. The _compute_fans() function in both init_ops.py and initializers_v2.py attempted to directly convert shape dimensions to int(), which fails for symbolic tensors with: TypeError: int() argument must be a string, a bytes-like object or a real number, not 'SymbolicTensor' Solution: --------- Modified _compute_fans() in both files to: 1. Use tensor_util.constant_value() to attempt extracting concrete values 2. Gracefully handle symbolic tensors with informative error messages 3. Provide clear guidance about using concrete shapes with XLA Changes: -------- 1. tensorflow/python/ops/init_ops.py - Added tensor_util import - Updated _compute_fans() with _to_int() helper function - Added informative error messages for dynamic shapes 2. tensorflow/python/keras/initializers/initializers_v2.py - Added tensor_util import - Updated _compute_fans() with same fix as init_ops.py - Ensures consistency across TF2 and Keras initializers 3. tensorflow/python/ops/test_xla_initializers_dynamic_shapes.py (new) - Comprehensive test suite validating the fix - Tests concrete shapes work with XLA - Tests dynamic shapes provide clear errors - Tests multiple initializer types 4. tensorflow/python/ops/demo_xla_initializers_fix.py (new) - Demonstration script showing the issue and solutions - Documents recommended patterns for XLA with initializers Testing: -------- The fix ensures: - Concrete shapes work correctly with XLA JIT compilation - Dynamic shapes fail with clear, actionable error messages - All variance scaling initializers (Glorot, He, Lecun) work properly - Backward compatibility is maintained for non-XLA code paths Workarounds for users: --------------------- 1. Use concrete shape values instead of tf.shape() 2. Initialize weights outside @tf.function(jit_compile=True) 3. Use tf.keras.layers.Dense with built-in initialization Fixes tensorflow#105334
- Detailed explanation of the problem and solution - Code examples showing what works and what doesn't - Testing instructions and expected outcomes - Impact analysis and next steps
…05333) Fixes tensorflow#105333 This commit fixes issue tensorflow#105333 where @tf.function(jit_compile=True) fails when returning dictionaries with mixed key types (e.g., strings and integers). Root Cause: ----------- When XLA JIT compilation flattens dictionaries, it sorts the keys to ensure deterministic ordering. However, Python 3 doesn't allow direct comparison between different types like int and str, causing: TypeError: '<' not supported between instances of 'int' and 'str' This occurred in _tf_core_sorted() and _tf_data_sorted() functions in nest_util.py when they called sorted() on dictionary keys containing mixed types. Solution: --------- Modified both sorting functions to use a fallback strategy: 1. First try direct sorting (for homogeneous key types) 2. If that fails, sort by (type_name, value) tuples 3. If that fails, sort by (type_name, str(value)) This ensures: - Deterministic ordering across all calls - Keys grouped by type (all ints, all strs, etc.) - Within each type group, sorted by value - Works with any combination of types Changes: -------- 1. tensorflow/python/util/nest_util.py - Updated _tf_core_sorted() with multi-level fallback sorting - Updated _tf_data_sorted() with same fix - Removed unhelpful error raising - Added clear comments explaining the strategy 2. tensorflow/python/util/test_mixed_dict_keys.py (new) - Comprehensive test suite validating the fix - Tests basic mixed keys (str + int) - Tests multiple type combinations - Tests nested dictionaries with mixed keys - Tests ordering consistency - Tests with and without XLA 3. tensorflow/python/util/demo_mixed_dict_keys.py (new) - Demonstration script showing the fix - Real-world multi-task model example - Documents sorting behavior - Validates consistency Testing: -------- The fix ensures: - Mixed-type dict keys work with XLA JIT compilation - Ordering is deterministic and consistent - Backward compatibility maintained for homogeneous keys - Works with nested structures Example: -------- # Now works with XLA: @tf.function(jit_compile=True) def mixed_keys(x): results = {} results['string_key'] = x results[123] = x + 1 return results Fixes tensorflow#105333
…eptive field calculation
CodersAcademy006
force-pushed
the
fix-xla-mixed-dict-keys
branch
from
a36f07b to
dd33a98
Compare
CodersAcademy006
force-pushed
the
fix-xla-mixed-dict-keys
branch
3 times, most recently
from
dd33a98 to
004b483
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.