Change extra arguments into pack_checkpoint function to save all parameters. Need for resume option.

Pull Request FAQ

• documentation
• contribution guide
• minimal examples section
• changelog for main framework updates
• Catalyst slack (#__questions channel) for issue discussion

Description

Related Issue

Type of Change

• Examples / docs / tutorials / contributors update
• Bug fix (non-breaking change which fixes an issue)
• Improvement (non-breaking change which improves an existing feature)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Checklist

• Have you updated tests for the new functionality?
• Have you added your new classes/functions to the docs?
• Have you updated the CHANGELOG?
• Have you run colab minimal CI/CD with latest and minimal requirements?
• Have you checked XLA integration with single and multiple processes?