def test_remat_scope_nested(self):
        """Test nested scopes with different rematerialization modes."""
        with RematScope(mode="full"):
            self.assertEqual(
                get_current_remat_mode().mode, "full"
            )  # Outer scope is "full"

            with RematScope(mode="activations"):
                self.assertEqual(
                    get_current_remat_mode().mode, "activations"
                )  # Inner scope is "activations"

            self.assertEqual(
                get_current_remat_mode().mode, "full"
            )  # Back to outer scope

        self.assertIsNone(
            get_current_remat_mode()
        )  # Mode is restored to None after all scopes
Test nested scopes with different rematerialization modes.
keras/src/backend/common/remat_test.py