inspect_ofa_ops.py

import sys
import inspect

# --- User Configuration for OFA Path ---
# Set this to the root of your cloned MIT OFA repository if the pip version doesn't work
# or if you want to be sure you're inspecting the repo version.
# If you believe the `pip install ofa` version is what's being used and is correct,
# you can leave OFA_REPO_PATH as None or a non-existent path.
OFA_REPO_PATH = None # Example: '/home/dgx-s-user2/once-for-all'
# OFA_REPO_PATH = '/home/dgx-s-user2/once-for-all' # <--- UNCOMMENT AND SET THIS if needed

if OFA_REPO_PATH and OFA_REPO_PATH not in sys.path:
    sys.path.insert(0, OFA_REPO_PATH)
    print(f"Added {OFA_REPO_PATH} to sys.path for OFA module inspection.")
# --- End User Configuration ---

def print_signature(cls, class_name_str):
    if cls is None:
        print(f"\n--- {class_name_str} not found/imported ---")
        return
    try:
        print(f"\n--- Signature for {class_name_str}.__init__ ---")
        sig = inspect.signature(cls.__init__)
        print(sig)
        print("\nParameters:")
        for param_name, param in sig.parameters.items():
            print(f"  {param_name}: {param.annotation} (default: {param.default})")
    except Exception as e:
        print(f"Could not get signature for {class_name_str}: {e}")

# Attempt to import the dynamic operations
DynamicConv2d_cls = None
DynamicBatchNorm2d_cls = None
DynamicSeparableConv2d_cls = None
DynamicLinear_cls = None
IdentityLayer_cls = None # From ofa.utils.layers if different from nn.Identity

try:
    from ofa.imagenet_classification.elastic_nn.modules.dynamic_op import (
        DynamicConv2d, DynamicBatchNorm2d, DynamicSeparableConv2d, DynamicLinear
    )
    DynamicConv2d_cls = DynamicConv2d
    DynamicBatchNorm2d_cls = DynamicBatchNorm2d
    DynamicSeparableConv2d_cls = DynamicSeparableConv2d
    DynamicLinear_cls = DynamicLinear
    print("Successfully imported dynamic_op modules.")
except ImportError as e_dyn:
    print(f"Failed to import from ofa.imagenet_classification.elastic_nn.modules.dynamic_op: {e_dyn}")
    print("Ensure OFA_REPO_PATH is correctly set if using a cloned repository, or 'ofa' pip package is complete.")

try:
    from ofa.utils.layers import IdentityLayer
    IdentityLayer_cls = IdentityLayer
    print("Successfully imported IdentityLayer from ofa.utils.layers.")
except ImportError as e_util:
    print(f"Failed to import IdentityLayer from ofa.utils.layers: {e_util}")


# Print signatures
print_signature(DynamicConv2d_cls, "DynamicConv2d")
print_signature(DynamicBatchNorm2d_cls, "DynamicBatchNorm2d")
print_signature(DynamicSeparableConv2d_cls, "DynamicSeparableConv2d")
print_signature(DynamicLinear_cls, "DynamicLinear")
print_signature(IdentityLayer_cls, "ofa.utils.layers.IdentityLayer")

# You can also check for active_in_channel, active_out_channel attributes if needed,
# though the __init__ signature is the immediate blocker.
# For example, if DynamicConv2d_cls:
#   instance = DynamicConv2d_cls(...) # with minimal valid args from its signature
#   print(hasattr(instance, 'active_in_channel'))
#   print(hasattr(instance, 'active_out_channel'))