Repository navigation
Fix Jacobian output dtypes for mixed precision differentiation - #129941
copybara-service[bot] merged 2 commits into
Conversation
google-ml-butler
Bot
added
size:M
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes dtype handling in tape.jacobian and tape.batch_jacobian when dealing with mixed-dtype sources and targets. Specifically, it ensures that the loop and output dtypes match the source dtypes rather than the target dtypes. It also adds comprehensive parameterized tests to verify mixed-dtype behavior for both jacobian and batch jacobian computations under various configurations. I have no feedback to provide as the changes are correct and well-tested.
keerthanakadiri
added
python
dmiltr3
left a comment
dmiltr3
left a comment
There was a problem hiding this comment.
Thanks for fixing mixed-precision dtype handling in GradientTape.jacobian and GradientTape.batch_jacobian. Please make two small updates before we merge this change:
- Normalize each source dtype with
dtypes.as_dtype(...)intensorflow/python/eager/backprop.py(lines 1195, 1341, and 1348).GradientTapesupports variable wrappers whose.dtypeattribute is a Python string, such askeras.VariableandVariableWithStringDtype. Insidepfor_ops.for_loop,TensorArray(dtype.base_dtype, iters)accesses.base_dtypedirectly, which raises anAttributeErroron string dtypes.
output = pfor_ops.for_loop(
loop_fn,
[dtypes.as_dtype(source.dtype) for source in flat_sources],
target_size,
parallel_iterations=parallel_iterations)And in batch_jacobian:
output = pfor_ops.for_loop(
loop_fn, dtypes.as_dtype(source.dtype), target_row_size,
parallel_iterations=parallel_iterations)Please also pass dtypes.as_dtype(source.dtype) to array_ops.zeros on line 1348 for consistency.
- Add a regression test in
tensorflow/python/eager/backprop_test.pyfor string-dtype variable wrappers. ReuseVariableWithStringDtype(fromtestUnconnectedGradientsVariableWithStringDtype) to testjacobianandbatch_jacobianwithexperimental_use_pfor=Falseand mixed dtypes.
github-project-automation
Bot
moved this from Assigned Reviewer
to Reviewer Requested Changes
in PR Queue
google-ml-butler
Bot
added
kokoro:force-run
github-project-automation
Bot
moved this from Reviewer Requested Changes
to Approved by Reviewer
in PR Queue
github-project-automation
Bot
moved this from Approved by Reviewer
to Merged
in PR Queue
Fixes #129940.
When a differentiable computation changes dtype, disabling pfor makes
GradientTape.jacobianandbatch_jacobianallocate TensorArrays with the target dtype, although the gradients have the source dtype. For example, a float32 source cast to float64 raises a TensorArray dtype error. Allocate each loop output with its corresponding source dtype, restoring the documented fallback for computations that cannot be vectorized.Use the source dtype for disconnected batch Jacobian zeros as well. This retains the existing zero-returning behavior in the default unconnected-gradient mode, while correcting the dtype to match connected Jacobians and
unconnected_gradients='zero'.Regression coverage includes nested sources with different dtypes, both float32/float64 conversion directions, eager and graph execution,
tf.function, both computation backends, and disconnected gradients in both unconnected-gradient modes.Validation on macOS arm64, Python 3.12.14, with tf-nightly 2.23.0.dev20261102:
JacobianTestandBatchJacobianTestclasses with the modifiedbackprop.pyloaded over the Nightly native runtime: 50 tests, 41 passed and 9 skipped under their existing decorators. No full Bazel build was run.backprop.py.tf.functionexecution. Both Jacobian implementations matched pfor results and the ordinary gradient dtype.git diff --checkpassed.