Fix TensorFlow Device Configuration After Import

[saksham-1304]

Fix #911

Fix TensorFlow Device Configuration After Import

Overview

This PR fixes an issue where importing tensorflow_quantum would trigger TensorFlow device initialization, preventing users from configuring devices afterward. This caused a RuntimeError: Physical devices cannot be modified after being initialized when attempting to set device configuration (e.g., GPU memory growth) after importing TFQ.

Problem Statement

Current Behavior (Before Fix)

import tensorflow as tf
import tensorflow_quantum as tfq  # Device initialization happens here!

gpus = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)  # RuntimeError!

Error: RuntimeError: Physical devices cannot be modified after being initialized

Root Cause

Multiple op modules in tensorflow_quantum called load_module() at module-level import time:

• tfq_utility_ops.py
• tfq_simulate_ops.py
• tfq_unitary_op.py
• tfq_adj_grad_op.py
• tfq_ps_util_ops.py
• noisy_samples_op.py
• noisy_sampled_expectation_op.py
• noisy_expectation_op.py

The load_module() function directly called load_library.load_op_library(), which initialized TensorFlow devices immediately.

Solution

Approach: Lazy Loading

Implemented a _LazyLoader class that defers the actual library loading until the ops are first accessed. This allows TensorFlow device configuration to happen after importing TFQ.

Key Changes

File: tensorflow_quantum/core/ops/load_module.py

1. Added _LazyLoader class:

   • Wraps the module loading process
   • Stores the module name but doesn't load until needed
   • Implements __getattr__() to load on first attribute access
   • Caches the loaded module for subsequent accesses

2. Modified load_module() function:

   • Returns a _LazyLoader instance instead of directly loading
   • Maintains the same interface for backward compatibility
   • Delegates all attribute access to the lazy loader

How It Works

# Before (Eager Loading)
UTILITY_OP_MODULE = load_module("_tfq_utility_ops.so")  # Loads immediately!

# After (Lazy Loading)
UTILITY_OP_MODULE = load_module("_tfq_utility_ops.so")  # Returns LazyLoader
UTILITY_OP_MODULE.tfq_append_circuit(...)              # Loads now!

Expected Behavior (After Fix)

import tensorflow as tf
import tensorflow_quantum as tfq  # No device initialization!

gpus = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)  # Works!

# Use TFQ normally
result = tfq.get_expectation_op()

Backward Compatibility

✅ Fully backward compatible

• The lazy loader implements __getattr__(), making it transparent to users
• All existing code continues to work without modifications
• No API changes or deprecations

Files Modified/Added

Modified

• tensorflow_quantum/core/ops/load_module.py - Core fix with lazy loading implementation

Added

• tensorflow_quantum/core/ops/load_module_test.py - Unit tests for lazy loader
• test_device_config_after_import.py - Integration test demonstrating the fix

Testing

Unit Tests (load_module_test.py)

- test_load_module_returns_lazy_loader()
- test_lazy_loader_defers_loading()
- test_lazy_loader_loads_on_attribute_access()
- test_lazy_loader_attribute_access_works()

Integration Test (test_device_config_after_import.py)

Demonstrates the fix in action:

# Can set GPU memory growth after importing TFQ
import tensorflow as tf
import tensorflow_quantum as tfq

gpus = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)  # ✓ Works!

Running Tests

# Run unit tests
python -m pytest tensorflow_quantum/core/ops/load_module_test.py -v

# Run integration test
python test_device_config_after_import.py

Performance Impact

Minimal impact:

• No overhead during TFQ import (faster import)
• Negligible overhead on first op usage (single function call to load)
• No overhead on subsequent op uses (cached module)

Use Cases Fixed

1. GPU Memory Configuration

import tensorflow as tf
import tensorflow_quantum as tfq

gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

2. Device Placement

import tensorflow as tf
import tensorflow_quantum as tfq

with tf.device('/GPU:0'):
    # Use TFQ ops here
    result = tfq.get_expectation_op()

3. Multiple GPU Setup

import tensorflow as tf
import tensorflow_quantum as tfq

# Configure GPUs first
tf.config.set_visible_devices(gpus[0], 'GPU')

# Then use TFQ
ops = tfq.get_sampling_op()

Breaking Changes

None. This is a non-breaking fix that improves behavior without changing the API.

Verification

The fix has been verified to:

• ✅ Allow device configuration after TFQ import
• ✅ Not break existing code
• ✅ Properly load ops on first use
• ✅ Cache loaded modules for efficiency
• ✅ Work with all TFQ op modules

Related Issues

• Fixes: RuntimeError: Physical devices cannot be modified after being initialized
• Allows users to follow best practices for GPU memory configuration
• Resolves unexpected behavior documented in issue

Documentation

For users encountering this issue, the recommended pattern is now:

import tensorflow as tf
import tensorflow_quantum as tfq

# Configure devices AFTER importing TFQ (now works!)
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError:
        pass

# Use TFQ normally
ops = tfq.get_expectation_op()

Implementation Details

_LazyLoader Class

• Location: tensorflow_quantum/core/ops/load_module.py
• Responsibility: Defer library loading until first attribute access
• Method: Implements __getattr__() to intercept attribute access
• Caching: Stores loaded module in _module attribute for reuse

Load Process

1. load_module("name.so") creates _LazyLoader instance
2. No loading happens at this point
3. When op is accessed (e.g., UTILITY_OP_MODULE.tfq_append_circuit):
   • __getattr__() is called
   • _load() method loads the actual library
   • Module is cached in _module
   • Attribute is retrieved and returned
4. Subsequent accesses use cached module

Rollback Plan

If needed, rollback is simple:

1. Revert changes to load_module.py
2. Remove test files
3. No other files need to be changed

Future Improvements

• Consider adding metrics to track lazy loading behavior
• Could extend pattern to other initialization-heavy imports
• Potential for similar lazy loading in other TensorFlow integration points

Questions?

For questions about this implementation, refer to:

• The inline comments in load_module.py
• The test cases in load_module_test.py
• The integration test in test_device_config_after_import.py

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.