p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] tests/test_checkpointer.py (python)
        # Compare parameters recursively, handling nested structure
        def compare_params(original, restored):
            if isinstance(original, dict):
                for k in original.keys():
                    compare_params(original[k], restored[k])
            else:
                self.assertTrue(jnp.allclose(original, restored))
        
        compare_params(state.params, restored["model"].params)