def scope(self):
        """Context manager to make the `Distribution` current."""
        original_scope = distribution()
        set_distribution(self)
        try:
            yield
        finally:
            set_distribution(original_scope)
Context manager to make the `Distribution` current.
keras/src/distribution/distribution_lib.py