Fix TensorFlow Device Configuration After Import

tensorflow_quantum/core/ops/load_module.py

@@ -21,26 +21,57 @@
 from tensorflow.python.platform import resource_loader
 
 
+class _LazyLoader:
+    """Lazily loads a TensorFlow op library on first attribute access.
+
+    This defers the call to `load_library.load_op_library` until the module
+    is actually used, preventing TensorFlow device initialization at import
+    time. This allows users to configure TensorFlow devices (e.g., enabling
+    memory growth) after importing tensorflow_quantum.
+    """
+
+    def __init__(self, name):
+        """Initialize the lazy loader.
+
+        Args:
+            name: The name of the module, e.g. "_tfq_simulate_ops.so"
+        """
+        self._name = name
+        self._module = None
+
+    def _load(self):
+        """Load the module if not already loaded."""
+        if self._module is None:
+            try:
+                path = resource_loader.get_path_to_datafile(self._name)
+                self._module = load_library.load_op_library(path)
+            except:
+                path = os.path.join(get_python_lib(),
+                                    "tensorflow_quantum/core/ops", self._name)
+                self._module = load_library.load_op_library(path)
+        return self._module
+
+    def __getattr__(self, name):
+        """Load the module on first attribute access and delegate."""
+        module = self._load()
+        return getattr(module, name)
+
+
 def load_module(name):
-    """Loads the module with the given name.
+    """Returns a lazy loader for the module with the given name.
 
-    First attempts to load the module as though it was embedded into the binary
-    using Bazel. If that fails, then it attempts to load the module as though
-    it was installed in site-packages via PIP.
+    The actual library loading is deferred until the module is first used.
+    This prevents TensorFlow device initialization at import time, allowing
+    users to configure TensorFlow devices after importing tensorflow_quantum.
 
     Args:
         name: The name of the module, e.g. "_tfq_simulate_ops.so"
 
     Returns:
-        A python module containing the Python wrappers for the Ops.
+        A lazy loader object that behaves like the loaded module but defers
+        loading until first attribute access.
 
     Raises:
-        RuntimeError: If the library cannot be found.
+        RuntimeError: If the library cannot be found when first accessed.
     """
-    try:
-        path = resource_loader.get_path_to_datafile(name)
-        return load_library.load_op_library(path)
-    except:
-        path = os.path.join(get_python_lib(), "tensorflow_quantum/core/ops",
-                            name)
-        return load_library.load_op_library(path)
+    return _LazyLoader(name)

tensorflow_quantum/core/ops/load_module_test.py

@@ -0,0 +1,51 @@
+# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for load_module lazy loading functionality."""
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow_quantum.core.ops.load_module import load_module, _LazyLoader
+
+
+class LoadModuleTest(tf.test.TestCase, parameterized.TestCase):
+    """Tests for load_module function and _LazyLoader class."""
+
+    def test_load_module_returns_lazy_loader(self):
+        """Test that load_module returns a _LazyLoader instance."""
+        loader = load_module("_tfq_utility_ops.so")
+        self.assertIsInstance(loader, _LazyLoader)
+
+    def test_lazy_loader_defers_loading(self):
+        """Test that _LazyLoader does not load the module on construction."""
+        loader = _LazyLoader("_tfq_utility_ops.so")
+        # _module should be None before any attribute access
+        self.assertIsNone(loader._module)
+
+    def test_lazy_loader_loads_on_attribute_access(self):
+        """Test that _LazyLoader loads the module on attribute access."""
+        loader = load_module("_tfq_utility_ops.so")
+        # Access an attribute to trigger loading
+        _ = loader.tfq_append_circuit
+        # Now _module should be loaded
+        self.assertIsNotNone(loader._module)
+
+    def test_lazy_loader_attribute_access_works(self):
+        """Test that attributes from the loaded module are accessible."""
+        loader = load_module("_tfq_utility_ops.so")
+        # Accessing an op should return a callable
+        self.assertTrue(callable(loader.tfq_append_circuit))
+
+
+if __name__ == '__main__':
+    tf.test.main()

test_device_config_after_import.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test that TensorFlow device configuration works after importing TFQ.
+
+This script verifies the fix for the issue where importing tensorflow_quantum
+before setting TensorFlow device configuration (e.g., enabling memory growth)
+resulted in a RuntimeError:
+
+    RuntimeError: Physical devices cannot be modified after being initialized
+
+The fix uses lazy loading of native op libraries to defer TensorFlow device
+initialization until the ops are actually used.
+
+Usage:
+    python test_device_config_after_import.py
+"""
+
+import sys
+import tensorflow as tf
+
+# Import tensorflow_quantum BEFORE configuring devices.
+# This used to trigger device initialization and cause errors.
+import tensorflow_quantum as tfq
+
+# Now try to configure devices - this should work without RuntimeError
+gpus = tf.config.list_physical_devices('GPU')
+if gpus:
+    try:
+        # Try setting memory growth - this would fail before the fix
+        tf.config.experimental.set_memory_growth(gpus[0], True)
+        print("SUCCESS: Device configuration after import works!")
+        print(f"  - Configured memory growth for GPU: {gpus[0]}")
+    except RuntimeError as e:
+        print(f"FAILED: {e}")
+        sys.exit(1)
+else:
+    # No GPU available, but we can still test that importing TFQ
+    # doesn't prematurely initialize devices by checking CPU config
+    cpus = tf.config.list_physical_devices('CPU')
+    print("SUCCESS: TFQ import did not prematurely initialize devices!")
+    print(f"  - Available CPUs: {cpus}")
+    print("  - No GPU available to test memory growth, but import test passed.")
+
+# Verify TFQ is actually usable after configuration
+print(f"  - TFQ version: {tfq.__version__}")