Fix torch keyword arguments in remat.

[hertschuh]

#21861

[gemini-code-assist]

Summary of Changes

Hello @hertschuh, 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 resolves an issue in the Keras PyTorch backend where the remat utility, used for gradient checkpointing, failed to correctly propagate keyword arguments to the torch.utils.checkpoint.checkpoint function. The change ensures that functions wrapped with remat can now properly receive and process keyword arguments, thereby enhancing the functionality and reliability of gradient checkpointing within the PyTorch backend. A new test case has been introduced to confirm this corrected behavior across various keyword argument scenarios.

Highlights

• Fix for remat keyword arguments: The remat function in the PyTorch backend has been updated to correctly pass keyword arguments to the underlying torch.utils.checkpoint.checkpoint function.
• New test case for remat with keyword arguments: A new test, test_remat_with_kwargs, has been added to remat_test.py to validate that remat properly handles functions that utilize keyword arguments, including explicit, default, and partial usage.

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 correctly fixes an issue in the PyTorch backend where remat would not forward keyword arguments. The change simply passes **kwargs to torch.utils.checkpoint.checkpoint. A new test, test_remat_with_kwargs, is added to verify the fix with various keyword argument scenarios, including full, partial, and default arguments. The implementation is correct and the test coverage is good.

[gemini-code-assist]

The test cases for different keyword argument scenarios are well-covered. For better maintainability and readability, you could consider parameterizing this test using subTest. This would make it more compact and easier to add new test cases in the future.

Suggested change

	# Test with keyword arguments
	remat_fn = backend.core.remat(fn_with_kwargs)
	result_with_kwargs = remat_fn(x, scale=2.0, offset=1.0)
	expected = fn_with_kwargs(x, scale=2.0, offset=1.0)
	self.assertAllClose(result_with_kwargs, expected)
	# Test with default keyword arguments
	result_with_defaults = remat_fn(x)
	expected_defaults = fn_with_kwargs(x)
	self.assertAllClose(result_with_defaults, expected_defaults)
	# Test with partial keyword arguments
	result_partial = remat_fn(x, scale=3.0)
	expected_partial = fn_with_kwargs(x, scale=3.0)
	self.assertAllClose(result_partial, expected_partial)
	remat_fn = backend.core.remat(fn_with_kwargs)
	test_cases = [
	("with_kwargs", {"scale": 2.0, "offset": 1.0}),
	("with_defaults", {}),
	("partial_kwargs", {"scale": 3.0}),
	]
	for name, kwargs in test_cases:
	with self.subTest(msg=name):
	result = remat_fn(x, **kwargs)
	expected = fn_with_kwargs(x, **kwargs)
	self.assertAllClose(result, expected)

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.57%. Comparing base (f2c00fe) to head (3459817).

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #21883   +/-   ##
=======================================
  Coverage   82.57%   82.57%           
=======================================
  Files         577      577           
  Lines       59599    59599           
  Branches     9351     9351           
=======================================
  Hits        49213    49213           
  Misses       7978     7978           
  Partials     2408     2408           

Flag	Coverage Δ
keras	82.39% <100.00%> (ø)
keras-jax	62.88% <0.00%> (ø)
keras-numpy	57.52% <0.00%> (ø)
keras-openvino	34.34% <0.00%> (ø)
keras-tensorflow	64.40% <0.00%> (ø)
keras-torch	63.57% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.