[BUG] Lazy import can break NAGL detection

[lohedges]

The BioSImSpace parameters code has auto-detection for openff-nagl support, which appears to be broken if any openff-toolkit object is imported after BioSimSpace is imported when using lazy loading, e.g.:

This works:

from openff.toolkit import Topology
import BioSimSpace as BSS

mol = BSS.Parameters.openff_unconstrained_2_2_0("C", use_nagl=True).getMolecule()

giving:

INFO:openff.nagl.nn._models:Using lookup table for property am1bcc_charges
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 0
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 1
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 2
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 3
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 4

But this fails, and falls back to AMBER AM1BCC:

import BioSimSpace as BSS
from openff.toolkit import Topology

mol = BSS.Parameters.openff_unconstrained_2_2_0("C", use_nagl=True).getMolecule()

giving:

INFO:openff.interchange.smirnoff._nonbonded:Charge section ToolkitAM1BCC, using charge method am1bcc, applied to topology atom index 0
INFO:openff.interchange.smirnoff._nonbonded:Charge section ToolkitAM1BCC, using charge method am1bcc, applied to topology atom index 1
INFO:openff.interchange.smirnoff._nonbonded:Charge section ToolkitAM1BCC, using charge method am1bcc, applied to topology atom index 2
INFO:openff.interchange.smirnoff._nonbonded:Charge section ToolkitAM1BCC, using charge method am1bcc, applied to topology atom index 3
INFO:openff.interchange.smirnoff._nonbonded:Charge section ToolkitAM1BCC, using charge method am1bcc, applied to topology atom index 4

However, running the script with the BSS_NO_LAZY_IMPORT environment variable set works, giving:

INFO:openff.nagl.nn._models:Using lookup table for property am1bcc_charges
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 0
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 1
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 2
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 3
INFO:openff.interchange.smirnoff._nonbonded:Preset charges applied to atom index 4

[lohedges]

The exception actual comes from importing openff.nag. This throws:

Failed to import openff.nagl: Only a single TORCH_LIBRARY can be used to register the namespace prims; please put all of your definitions in a single TORCH_LIBRARY block.  If you were trying to specify implementations, consider using TORCH_LIBRARY_IMPL (which can be duplicated).  If you really intended to define operators for a single namespace in a distributed way, you can use TORCH_LIBRARY_FRAGMENT to explicitly indicate this.  Previous registration of TORCH_LIBRARY was registered at /dev/null:488; latest registration was registered at /dev/null:488

[lohedges]

This is weird. The following fails at the torch import, but the exception isn't even raised (presumably it crashes somehow).

import BioSimSpace as BSS
from openff.toolkit import Topology

print("HERE")
try:
    import torch
except ImportError:
    raise("Failed to import torch") from e
print("THERE")

Giving:

INFO:numexpr.utils:Note: NumExpr detected 20 cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 16.
INFO:numexpr.utils:NumExpr defaulting to 16 threads.
HERE
INFO:torch._subclasses.fake_tensor:FakeTensor cache stats:
INFO:torch._subclasses.fake_tensor:  cache_hits: 0
INFO:torch._subclasses.fake_tensor:  cache_misses: 0

If I reorder the openff and BioSimSpace imports, or run without lazy importing BioSimSpace, then I see:

INFO:numexpr.utils:Note: NumExpr detected 20 cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 16.
INFO:numexpr.utils:NumExpr defaulting to 16 threads.
HERE
THERE

The same thing happens when replacing the openff import with import torch, so it's definitely a torch issue. It doesn't look like there are any direct imports of torch within BioSimSpace or Sire, at least anything that would be loaded directly at startup.