Upload 597 files

.gitattributes

@@ -36,3 +36,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 RFdiffusion/env/SE3Transformer/images/se3-transformer.png filter=lfs diff=lfs merge=lfs -text
 RFdiffusion/img/diffusion_protein_gradient_2.jpg filter=lfs diff=lfs merge=lfs -text
 RFdiffusion/pyrosetta-2023.14+release.7132bdc754a-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
+env/SE3Transformer/images/se3-transformer.png filter=lfs diff=lfs merge=lfs -text
+img/diffusion_protein_gradient_2.jpg filter=lfs diff=lfs merge=lfs -text

.github/CODEOWNERS

@@ -0,0 +1,2 @@
+# Benchmark scripts
+/.rosetta-ci @lyskov

.rosetta-ci/.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+results/
+benchmark.ubuntu.ini

.rosetta-ci/benchmark.py

@@ -0,0 +1,410 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   benchmark.py
+## @brief  Run arbitrary Rosetta testing script
+## @author Sergey Lyskov
+
+from __future__ import print_function
+
+import os, os.path, sys, shutil, json, platform, re
+import codecs
+
+from importlib.machinery import SourceFileLoader
+
+from configparser import ConfigParser, ExtendedInterpolation
+import argparse
+
+from tests import *  # execute, Tests states and key names
+from hpc_drivers import *
+
+
+# Calculating value of Platform dict
+Platform = {}
+if sys.platform.startswith("linux"):
+    Platform['os'] = 'ubuntu' if os.path.isfile('/etc/lsb-release') and 'Ubuntu' in open('/etc/lsb-release').read() else 'linux'  # can be linux1, linux2, etc
+elif sys.platform == "darwin" :      Platform['os'] = 'mac'
+elif sys.platform == "cygwin" :      Platform['os'] = 'cygwin'
+elif sys.platform == "win32" :       Platform['os'] = 'windows'
+else:                                Platform['os'] = 'unknown'
+
+#Platform['arch'] = platform.architecture()[0][:2]  # PlatformBits
+Platform['compiler'] = 'gcc' if Platform['os'] == 'linux' else 'clang'
+
+Platform['python'] = sys.executable
+
+
+def load_python_source_from_file(module_name, module_path):
+    ''' replacment for deprecated imp.load_source
+    '''
+    return SourceFileLoader(module_name, module_path).load_module()
+
+
+class Setup(object):
+    __slots__ = 'test working_dir platform config compare debug'.split()  # version daemon path_to_previous_test
+    def __init__(self, **attrs):
+        #self.daemon = True
+        for k, v in attrs.items():
+            if k in self.__slots__: setattr(self, k, v)
+
+
+def setup_from_options(options):
+    ''' Create Setup object based on user supplied options, config files and auto-detection
+    '''
+    platform = dict(Platform)
+
+    if options.suffix: options.suffix = '.' + options.suffix
+
+    platform['extras'] = options.extras.split(',') if options.extras else []
+    platform['python'] = options.python
+    #platform['options'] = json.loads( options.options ) if options.options else {}
+
+    if options.memory: memory = options.memory
+    elif platform['os'] in ['linux', 'ubuntu']: memory = int( execute('Getting memory info...', 'free -m', terminate_on_failure=False, silent=True, silence_output_on_errors=True, return_='output').split('\n')[1].split()[1]) // 1024
+    elif platform['os'] == 'mac':   memory = int( execute('Getting memory info...', 'sysctl -a | grep hw.memsize', terminate_on_failure=False, silent=True, silence_output_on_errors=True, return_='output').split()[1]) // 1024 // 1024 // 1024
+
+    platform['compiler'] = options.compiler
+
+    if os.path.isfile(options.config):
+        with open(options.config) as f:
+            if '%(here)s' in f.read():
+                print(f"\n\n>>> ERROR file `{options.config}` seems to be in outdated format! Please use benchmark.template.ini to update it.")
+                sys.exit(1)
+
+        user_config = ConfigParser(
+            dict(
+                _here_ = os.path.abspath('./'),
+                _user_home_ = os.environ['HOME']
+            ),
+            interpolation = ExtendedInterpolation()
+        )
+
+        with open(options.config) as f: user_config.readfp(f)
+
+    else:
+        print(f"\n\n>>> Config file `{options.config}` not found. You may want to manually copy `benchmark.ini.template` to `{options.config}` and edit the settings\n\n")
+        user_config = ConfigParser()
+        user_config.set('main', 'cpu_count',  '1')
+        user_config.set('main', 'hpc_driver', 'MultiCore')
+        user_config.set('main', 'branch',     'unknown')
+        user_config.set('main', 'revision',   '42')
+        user_config.set('main', 'user_name',  'Jane Roe')
+        user_config.set('main', 'user_email', 'jane.roe@university.edu')
+        user_config.add_section('main')
+
+    if options.jobs: user_config.set('main', 'cpu_count', str(options.jobs) )
+    user_config.set('main', 'memory',    str(memory) )
+
+    if options.mount:
+        for m in options.mount:
+            key, _, path = m.partition(':')
+            user_config.set('mount', key, path)
+
+    #config = Config.items('config')
+    #for section in config.sections(): print('Config section: ', section, dict(config.items(section)))
+    #config = { section: dict(Config.items(section)) for section in Config.sections() }
+
+    config = { k : d for k, d in user_config['main'].items() if k not in user_config[user_config.default_section] }
+    config['mounts'] = { k : d for k, d in user_config['mount'].items() if k not in user_config[user_config.default_section] }
+
+    #print(json.dumps(config, sort_keys=True, indent=2)); sys.exit(1)
+
+    #config.update( config.pop('config').items() )
+
+    config = dict(config,
+                  cpu_count = user_config.getint('main', 'cpu_count'),
+                  memory = memory,
+                  revision = user_config.getint('main', 'revision'),
+                  emulation=True,
+    )  # debug=options.debug,
+
+    if 'results_root' not in config: config['results_root'] = os.path.abspath('./results/')
+
+    if 'prefix' in config:
+        assert os.path.isabs( config['prefix'] ), f'ERROR: `prefix` path must be absolute! Got: {config["prefix"]}'
+
+    else: config['prefix'] = os.path.abspath( config['results_root'] + '/prefix')
+
+    config['merge_head'] = options.merge_head
+    config['merge_base'] = options.merge_base
+
+    if options.skip_compile is not None: config['skip_compile'] = options.skip_compile
+
+    #print(f'Results path: {config["results_root"]}')
+    #print('Config:{}, Platform:{}'.format(json.dumps(config, sort_keys=True, indent=2), Platform))
+
+    if options.compare: print('Comparing tests {} with suffixes: {}'.format(options.args, options.compare) )
+    else: print('Running tests: {}'.format(options.args) )
+
+    if len(options.args) != 1: print('Error: Single test-name-to-run should be supplied!');  sys.exit(1)
+    else:
+        test = options.args[0]
+        if test.startswith('tests/'): test = test.partition('tests/')[2][:-3]  # removing dir prefix and .py suffix
+
+    if options.compare:
+        compare = options.compare[0], options.compare[1]  # (this test suffix, previous test suffix)
+        working_dir = os.path.abspath( config['results_root'] + f'/{platform["os"]}.{test}' )  # will be a root dir with sub-dirs (options.compare[0], options.compare[1])
+    else:
+        compare = None
+        working_dir = os.path.abspath( config['results_root'] + f'/{platform["os"]}.{test}{options.suffix}' )
+
+
+    if os.path.isdir(working_dir): shutil.rmtree(working_dir);  #print('Removing old job dir %s...' % working_dir)  # remove old dir if any
+    os.makedirs(working_dir)
+
+    setup = Setup(
+        test        = test,
+        working_dir = working_dir,
+        platform    = platform,
+        config      = config,
+        compare     = compare,
+        debug       = options.debug,
+        #daemon      = False,
+    )
+
+    setup_as_json = json.dumps( { k : getattr(setup, k) for k in setup.__slots__}, sort_keys=True, indent=2)
+    with open(working_dir + '/.setup.json', 'w') as f: f.write(setup_as_json)
+
+    #print(f'Detected hardware platform: {Platform}')
+    print(f'Setup: {setup_as_json}')
+    return setup
+
+
+def truncate_log(log):
+    _max_log_size_  = 1024*1024*1
+    _max_line_size_ = _max_log_size_ // 2
+
+    if len(log) > _max_log_size_:
+        new = log
+        lines = log.split('\n')
+
+        if len(lines) > 256:
+            new_lines = lines[:32] + ['...truncated...'] + lines[-128:]
+            new = '\n'.join(new_lines)
+
+        if len(new) > _max_log_size_: # special case for Ninja logs that does not use \n
+            lines = re.split(r'[\r\n]*', log)  #t.log.split('\r')
+            if len(lines) > 256: new = '\n'.join( lines[:32] + ['...truncated...'] + lines[-128:] )
+
+        if len(new) > _max_log_size_: # going to try to truncate each individual line...
+            print(f'Trying to truncate log line-by-line...')
+            new = '\n'.join( (
+                ( line[:_max_line_size_//3] + '...truncated...' + line[-_max_line_size_//3:] ) if line > _max_line_size_ else line
+                for line in new_lines ) )
+
+        if len(new) > _max_log_size_: # fall-back strategy in case all of the above failed...
+            print(f'WARNING: could not truncate log line-by-line, falling back to raw truncate...')
+            new = 'WARNING: could not truncate test log line-by-line, falling back to raw truncate!\n...truncated...\n' + ( '\n'.join(lines) )[-_max_log_size_+256:]
+
+        print( 'Trunacting test output log: {0}MiB --> {1}MiB'.format(len(log)/1024/1024, len(new)/1024/1024) )
+
+        log = new
+
+    return log
+
+def truncate_results_logs(results):
+    results[_LogKey_] = truncate_log( results[_LogKey_] )
+    if _ResultsKey_ in results  and  _TestsKey_ in results[_ResultsKey_]:
+        tests = results[_ResultsKey_][_TestsKey_]
+        for test in tests:
+            tests[test][_LogKey_] = truncate_log( tests[test][_LogKey_] )
+
+
+def find_test_description(test_name, test_script_file_name):
+    ''' return content of test-description file if any or None if no description was found
+    '''
+
+    def find_description_file(prefix, test_name):
+        fname = prefix + test_name + '.md'
+        if os.path.isfile(fname): return fname
+        return prefix + 'md'
+
+    description_file_name =  find_description_file( test_script_file_name[:-len('command.py')] + 'description.', test_name) if test_script_file_name.endswith('/command.py') else find_description_file(test_script_file_name[:-len('py')], test_name)
+
+    if description_file_name  and  os.path.isfile(description_file_name):
+        print(f'Found test suite description in file: {description_file_name!r}')
+        with open(description_file_name, encoding='utf-8', errors='backslashreplace') as f: description = f.read()
+        return description
+
+    else: return None
+
+
+
+def run_test(setup):
+    #print(f'{setup!r}')
+    suite, rest = setup.test.split('.'), []
+    while suite:
+        #print( f'suite: {suite}, test: {rest}' )
+
+        file_name = '/'.join( ['tests'] + suite ) + '.py'
+        if os.path.isfile(file_name): break
+
+        file_name = '/'.join( ['tests'] + suite ) + '/command.py'
+        if os.path.isfile(file_name): break
+
+        rest.insert(0, suite.pop())
+
+
+    test = '.'.join( suite + rest )
+    test_name = '.'.join(rest)
+
+    print( f'Loading test from: {file_name}, suite+test: {test!r}, test: {test_name!r}' )
+    #test_suite = imp.load_source('test_suite', file_name)
+    test_suite = load_python_source_from_file('test_suite', file_name)
+
+    test_description = find_test_description(test_name, file_name)
+
+    if setup.compare:
+        #working_dir_1 = os.path.abspath( config['results_root'] + f'/{Platform["os"]}.{test}.{Options.compare[0]}' )
+        working_dir_1 = setup.working_dir + f'/{setup.compare[0]}'
+
+        working_dir_2        = setup.compare[1]  and  ( setup.working_dir + f'/{setup.compare[1]}' )
+        res_2_json_file_path = setup.compare[1]  and  f'{working_dir_2}/.execution.results.json'
+
+        with open(working_dir_1 + '/.execution.results.json') as f: res_1 = json.load(f).get(_ResultsKey_)
+
+        if setup.compare[1] and ( not os.path.isfile(res_2_json_file_path) ):
+            setup.compare[1] = None
+            state_override = _S_failed_
+        else:
+            state_override = None
+
+        if setup.compare[1] == None: res_2, working_dir_2 = None, None
+        else:
+            with open(res_2_json_file_path) as f: res_2 = json.load(f).get(_ResultsKey_)
+
+        res = test_suite.compare(test, res_1, working_dir_1, res_2, working_dir_2)
+
+        if state_override:
+            log_prefix = \
+                f'WARNING: Previous test results does not have `.execution.results.json` file, so comparision with None was performed instead!\n' \
+                f'WARNING: Overriding calcualted test state `{res[_StateKey_]}` → `{_S_failed_}`...\n\n'
+
+            res[_LogKey_] = log_prefix + res[_LogKey_]
+            res[_StateKey_] = _S_failed_
+
+
+        # # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages.
+        # with codecs.open(setup.working_dir+'/.comparison.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write( truncate_log( res[_LogKey_] ) )
+        # res[_LogKey_] = truncate_log( res[_LogKey_] )
+
+        # # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages.
+        with codecs.open(setup.working_dir+'/.comparison.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write(res[_LogKey_])
+        truncate_results_logs(res)
+
+        print( 'Comparison finished with output:\n{}'.format( res[_LogKey_] ) )
+
+        with open(setup.working_dir+'/.comparison.results.json', 'w') as f: json.dump(res, f, sort_keys=True, indent=2)
+
+        #print( 'Comparison finished with results:\n{}'.format( json.dumps(res, sort_keys=True, indent=2) ) )
+        if 'summary' in res: print('Summary section:\n{}'.format( json.dumps(res['summary'], sort_keys=True, indent=2) ) )
+
+        print( f'Output results of this comparison saved to {working_dir_1}/.comparison.results.json\nComparison log saved into {working_dir_1}/.comparison.log.txt' )
+
+
+    else:
+        working_dir = setup.working_dir  #os.path.abspath( setup.config['results_root'] + f'/{platform["os"]}.{test}{options.suffix}' )
+
+        hpc_driver_name = setup.config['hpc_driver']
+        hpc_driver = None if hpc_driver_name in ['', 'none'] else eval(hpc_driver_name + '_HPC_Driver')(working_dir, setup.config, tracer=print, set_daemon_message=lambda x:None)
+
+        api_version = test_suite._api_version_ if hasattr(test_suite, '_api_version_') else ''
+
+        # if api_version < '1.0':
+        #     res = test_suite.run(test=test_name, rosetta_dir=os.path.abspath('../..'), working_dir=working_dir, platform=dict(Platform), jobs=Config.cpu_count, verbose=True, debug=Options.debug)
+        # else:
+
+        if api_version == '1.0': res = test_suite.run(test=test_name, repository_root=os.path.abspath('./..'), working_dir=working_dir, platform=dict(setup.platform), config=setup.config, hpc_driver=hpc_driver, verbose=True, debug=setup.debug)
+        else:
+            print(f'Test benchmark api_version={api_version} is not supported!'); sys.exit(1)
+
+        if not isinstance(res, dict): print(f'Test returned result of type {type(res)} while dict-like object was expected, please check that test-script have correct `return` statment! Terminating...'); sys.exit(1)
+
+        # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages
+        with codecs.open(working_dir+'/.execution.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write( res[_LogKey_] )
+
+        # res[_LogKey_] = truncate_log( res[_LogKey_] )
+        truncate_results_logs(res)
+
+        if _DescriptionKey_ not in res: res[_DescriptionKey_] = test_description
+
+        if res[_StateKey_] not in _S_Values_: print( 'Warning!!! Test {} failed with unknow result code: {}'.format(test_name, res[_StateKey_]) )
+        else: print( f'Test {test} finished with output:\n{res[_LogKey_]}\n----------------------------------------------------------------\nState: {res[_StateKey_]!r} | ', end='')
+
+        # JSON by default serializes to an ascii-encoded format
+        with open(working_dir+'/.execution.results.json', 'w') as f: json.dump(res, f, sort_keys=True, indent=2)
+
+        print( f'Output and full log of this test saved to:\n{working_dir}/.execution.results.json\n{working_dir}/.execution.log.txt' )
+
+
+
+
+
+
+def main(args):
+    ''' Script to Run arbitrary Rosetta test
+    '''
+    parser = argparse.ArgumentParser(usage="Main testing script to run tests in the tests directory. "
+                                           "Use the --skip-compile to skip the build phase when testing locally. "
+                                           "Example Command: /benchmark.py -j2 integration.valgrind")
+
+    parser.add_argument('-j', '--jobs', default=0, type=int, help="Number of processors to use on when building. (default: use value from config file or 1)")
+
+    parser.add_argument('-m', '--memory', default=0, type=int, help="Amount of memory to use (default: use 2Gb per job")
+
+    parser.add_argument('--compiler', default=Platform['compiler'], help="Compiler to use")
+
+    #parser.add_argument('--python', default=('3.9' if Platform['os'] == 'mac' else '3.6'), help="Python interpreter to use")
+    parser.add_argument('--python', default=f'{sys.version_info.major}.{sys.version_info.minor}.s', help="Specify version of Python interpreter to use, for example '3.9'. If '.s' added to end of version string then use the same interpreter that was used to start this script. Default: '?.?.s'")
+
+    parser.add_argument("--extras", default='', help="Specify scons extras separated by ',': like --extras=mpi,static" )
+
+    parser.add_argument("--debug", action="store_true", dest="debug", default=False, help="Run specified test in debug mode (not with debug build!) this mean different things and depend on the test. Could be: skip the build phase, skip some of the test phases and so on. [off by default]" )
+
+    parser.add_argument("--suffix", default='', help="Specify ending suffix for test output dir. This is useful when you want to save test results in different dir for later comparison." )
+
+    parser.add_argument("--compare", nargs=2, help="Do not run the tests but instead compare previous results. Use --compare suffix1 suffix2" )
+
+    parser.add_argument("--config", default='benchmark.{os}.ini'.format(os=Platform['os']), action="store", help="Location of .ini file with additional options configuration. Optional.")
+
+    parser.add_argument("--skip-compile", dest='skip_compile', default=None, action="store_true", help="Skip the compilation phase. Assumes the binaries are already compiled locally.")
+
+    #parser.add_argument("--results-root", default=None, action="store", help="Location of `results` dir default is to use `./results`")
+
+    parser.add_argument("--setup", default=None, help="Specify JSON file with setup information. When this option supplied all other config and commandline options is ignored and auto-detection disable. Test, platform info will be gathered from provided JSON file. This option is designed to be used in daemon mode." )
+
+    parser.add_argument("--merge-head", default='HEAD', help="Specify SHA1/branch-name that will be used for `merge-head` value when simulating PR testing" )
+
+    parser.add_argument("--merge-base", default='origin/master', help="Specify SHA1/branch-name that will be used for `merge-base` value when simulating PR testing" )
+
+    parser.add_argument("--mount", action="append", help="Specify one of the mount points, like: --mount release_root:/some/path. This option could be used multiple times if needed" )
+
+
+    parser.add_argument('args', nargs=argparse.REMAINDER)
+
+    options = parser.parse_args(args=args[1:])
+
+    if any( [a.startswith('-') for a in options.args] ) :
+        print( '\nWARNING WARNING WARNING WARNING\n' )
+        print( '\tInterpreting', ' '.join(["'"+a+"'" for a in options.args if a.startswith('-')]), 'as test name(s), rather than as option(s).'  )
+        print( "\tTry moving it before any test name, if that's not what you want."  )
+        print( '\nWARNING WARNING WARNING WARNING\n'  )
+
+
+    if options.setup:
+        with open(options.setup) as f: setup = Setup( **json.load(f) )
+
+    else:
+        setup = setup_from_options(options)
+
+    run_test(setup)
+
+
+if __name__ == "__main__": main(sys.argv)

.rosetta-ci/benchmark.template.ini

@@ -0,0 +1,40 @@
+#
+# Benchmark script configuration file. Some of the tests require some system specific options to run. Please see benchmark.ini.template for list of available options.
+#
+
+[DEFAULT]
+
+[main] # additional config-options for various tests. All this fields will be pass as keys in 'config' function argument
+
+# how many jobs daemon can run on host machine (this is not related to HPC jobs)
+cpu_count = 24
+
+# how many memory in GB daemon can use on host machine (approximation, float)
+memory = 64
+
+# user name and email for user who submitted this test
+user_name  = Jane Roe
+user_email = jane.roe@university.edu
+
+# HPC Driver, might have one of the following values: MultiCore, Condor, Slurm or none if no HPC Driver should be configured
+hpc_driver = MultiCore
+
+# when running by daemons branch:revision will be set to appropriate values to represent currently checked version of main repository
+branch = unknown
+revision = 42
+
+# path to directory where test results will be stored
+results_root = ${_here_}/results
+
+release_root = ./results/_release_
+
+[slurm]
+# head-node host name, if specified will be used to submit jobs
+head_node =
+
+
+[mount]
+# list of key:path pairs that will be avalible as config.mounts during test run
+
+# path to releases, leave empty if release production should not be supported by this daemon
+release_root = ${_here_}/release

.rosetta-ci/hpc_drivers/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+from .multicore import MultiCore_HPC_Driver
+from .slurm     import Slurm_HPC_Driver

.rosetta-ci/hpc_drivers/base.py

@@ -0,0 +1,210 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import os, sys, subprocess, stat
+import time as time_module
+import signal as signal_module
+
+class NT:  # named tuple
+    def __init__(self, **entries): self.__dict__.update(entries)
+    def __repr__(self):
+        r = 'NT: |'
+        for i in dir(self):
+            if not i.startswith('__') and not isinstance(getattr(self, i), types.MethodType): r += '{} --> {}, '.format(i, getattr(self, i))
+        return r[:-2]+'|'
+
+
+
+class HPC_Exception(Exception):
+    def __init__(self, value): self.value = value
+    def __str__(self): return self.value
+
+
+
+def execute(message, command_line, return_='status', until_successes=False, terminate_on_failure=True, silent=False, silence_output=False, tracer=print):
+    if not silent: tracer(message);  tracer(command_line); sys.stdout.flush();
+    while True:
+
+        p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        output, errors = p.communicate()
+
+        output = output + errors
+
+        output = output.decode(encoding="utf-8", errors="replace")
+
+        exit_code = p.returncode
+
+        if exit_code  and  not (silent or silence_output): tracer(output); sys.stdout.flush();
+
+        if exit_code and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        tracer( "Error while executing {}: {}\n".format(message, output) )
+        tracer("Sleeping 60s... then I will retry...")
+        sys.stdout.flush();
+        time.sleep(60)
+
+    if return_ == 'tuple': return(exit_code, output)
+
+    if exit_code and terminate_on_failure:
+        tracer("\nEncounter error while executing: " + command_line)
+        if return_==True: return True
+        else: print("\nEncounter error while executing: " + command_line + '\n' + output); sys.exit(1)
+
+    if return_ == 'output': return output
+    else: return False
+
+
+def Sleep(time_, message, dict_={}):
+    ''' Fancy sleep function '''
+    len_ = 0
+    for i in range(time_, 0, -1):
+        #print "Waiting for a new revision:%s... Sleeping...%d     \r" % (sc.revision, i),
+        msg = message.format( **dict(dict_, time_left=i) )
+        print( msg, end='' )
+        len_ = max(len_, len(msg))
+        sys.stdout.flush()
+        time_module.sleep(1)
+
+    print( ' '*len_ + '\r',  end='' ) # erazing sleep message
+
+
+# Abstract class for HPC job submission
+class HPC_Driver:
+    def __init__(self, working_dir, config, tracer=lambda x:None, set_daemon_message=lambda x:None):
+        self.working_dir = working_dir
+        self.config = config
+        self.cpu_usage  = 0.0  # cummulative cpu usage in hours
+        self.tracer     = tracer
+        self.set_daemon_message = set_daemon_message
+
+        self.cpu_count = self.config['cpu_count'] if type(config) == dict else self.config.getint('DEFAULT', 'cpu_count')
+
+        self.jobs = []  # list of all jobs currently running by this driver, Job class is driver depended, could be just int or something more complex
+
+        self.install_signal_handler()
+
+
+    def __del__(self):
+        self.remove_signal_handler()
+
+
+    def execute(self, executable, arguments, working_dir, log_dir=None, name='_no_name_', memory=256, time=24, shell_wrapper=False, block=True):
+        ''' Execute given command line on HPC cluster, must accumulate cpu hours in self.cpu_usage '''
+        if log_dir==None: log_dir=self.working_dir
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + '/hpc.{}.shell_wrapper.sh'.format(name))
+            with file(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        return self.submit_serial_hpc_job(name=name, executable=executable, arguments=arguments, working_dir=working_dir, log_dir=log_dir, jobs_to_queue=1, memory=memory, time=time, block=block, shell_wrapper=shell_wrapper)
+
+
+
+    @property
+    def number_of_cpu_per_node(self):
+        must_be_implemented_in_inherited_classes
+
+    @property
+    def maximum_number_of_mpi_cpu(self):
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, memory=512, time=12, block=True, process_coefficient="1", requested_nodes=1, requested_processes_per_node=1):
+        ''' submit jobs as MPI job
+            process_coefficient should be string representing fraction of process to launch on each node, for example '3 / 4' will start only 75% of MPI process's on each node
+        '''
+        must_be_implemented_in_inherited_classes
+
+
+    def cancel_all_jobs(self):
+        ''' Cancel all HPC jobs known to this driver, use this as signal handler for script termination '''
+        for j in self.jobs: self.cancel_job(j)
+
+    def block_until(self, silent, fn, *args, **kwargs):
+        '''
+        **fn must have the driver as the first argument**
+        example:
+        def fn(driver):
+            jobs = list(driver.jobs)
+            jobs = [job for job in jobs if not driver.complete(job)]
+            if len(jobs) <= 8:
+                return False # stops sleeping
+            return True # continues sleeping
+
+        for x in range(100):
+            hpc_driver.submit_hpc_job(...)
+            hpc_driver.block_until(False, fn)
+        '''
+        while fn(self, *args, **kwargs):
+            sys.stdout.flush()
+            time_module.sleep(60)
+            if not silent:
+                Sleep(1, '"Waiting for HPC job(s) to finish, sleeping {time_left}s\r')
+
+    def wait_until_complete(self, jobs=None, callback=None, silent=False):
+        ''' Helper function, wait until given jobs list is finished, if no argument is given waits until all jobs known by driver is finished '''
+        jobs = jobs if jobs else self.jobs
+
+        while jobs:
+            for j in jobs[:]:
+                if self.complete(j): jobs.remove(j)
+
+            if jobs:
+                #total_cpu_queued  = sum( [j.jobs_queued  for j in jobs] )
+                #total_cpu_running = sum( [j.cpu_running for j in jobs] )
+                #self.set_daemon_message("Waiting for HPC job(s) to finish... [{} process(es) in queue, {} process(es) running]".format(total_cpu_queued, total_cpu_running) )
+                #self.tracer("Waiting for HPC job(s) [{} process(es) in queue, {} process(es) running]...  \r".format(total_cpu_queued, total_cpu_running), end='')
+                #print "Waiting for {} HPC jobs to finish... [{} jobs in queue, {} jobs running]... Sleeping 32s...     \r".format(total_cpu_queued, cpu_queued+cpu_running, cpu_running),
+
+                self.set_daemon_message("Waiting for HPC {} job(s) to finish...".format( len(jobs) ) )
+                #self.tracer("Waiting for HPC {} job(s) to finish...".format( len(jobs) ) )
+
+                sys.stdout.flush()
+
+                if callback: callback()
+
+                if silent: time_module.sleep(64*1)
+                else: Sleep(64, '"Waiting for HPC {n_jobs} job(s) to finish, sleeping {time_left}s    \r', dict(n_jobs=len(jobs)))
+
+
+
+    _signals_ = [signal_module.SIGINT, signal_module.SIGTERM, signal_module.SIGABRT]
+    def install_signal_handler(self):
+        def signal_handler(signal_, frame):
+            self.tracer('Recieved signal:{}... Canceling HPC jobs...'.format(signal_) )
+            self.cancel_all_jobs()
+            self.set_daemon_message( 'Remote daemon got terminated with signal:{}'.format(signal_) )
+            sys.exit(1)
+
+        for s in self._signals_: signal_module.signal(s, signal_handler)
+
+
+    def remove_signal_handler(self):  # do we really need this???
+        try:
+            for s in self._signals_: signal_module.signal(s, signal_module.SIG_DFL)
+            #print('remove_signal_handler: done!')
+
+        except TypeError:
+            #print('remove_signal_handler: interpreted terminating, skipping remove_signal_handler...')
+            pass
+
+
+    def cancel_job(self, job_id):
+        must_be_implemented_in_inherited_classes
+
+
+    def complete(self, job_id):
+        ''' Return job completion status. Return True if job complered and False otherwise
+        '''
+        must_be_implemented_in_inherited_classes

.rosetta-ci/hpc_drivers/multicore.py

@@ -0,0 +1,184 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import time as time_module
+import codecs
+import signal
+
+import os, sys
+
+try:
+    from .base import *
+
+except ImportError: # workaround for B2 back-end's
+    import imp
+    imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) + '/base.py')  # A bit of Python magic here, what we trying to say is this: from base import *, but path to base is calculated from our source location  # from base import HPC_Driver, execute, NT
+
+
+class MultiCore_HPC_Driver(HPC_Driver):
+
+    class JobID:
+        def __init__(self, pids=None):
+            self.pids = pids if pids else []
+
+
+        def __bool__(self): return bool(self.pids)
+
+
+        def __len__(self): return len(self.pids)
+
+
+        def add_pid(self, pid): self.pids.append(pid)
+
+
+        def remove_completed_pids(self):
+            for pid in self.pids[:]:
+                try:
+                    r = os.waitpid(pid, os.WNOHANG)
+                    if r == (pid, 0): self.pids.remove(pid)  # process have ended without error
+                    elif r[0] == pid :  # process ended but with error, special case we will have to wait for all process to terminate and call system exit.
+                        #self.cancel_job()
+                        #sys.exit(1)
+                        self.pids.remove(pid)
+                        print('ERROR: Some of the HPC jobs terminated abnormally! Please see HPC logs for details.')
+
+                except ChildProcessError: self.pids.remove(pid)
+
+
+        def cancel(self):
+            for pid in self.pids:
+                try:
+                    os.killpg(os.getpgid(pid), signal.SIGKILL)
+                except ChildProcessError: pass
+
+            self.pids = []
+
+
+
+    def __init__(self, *args, **kwds):
+        HPC_Driver.__init__(self, *args, **kwds)
+        #print(f'MultiCore_HPC_Driver: cpu_count: {self.cpu_count}')
+
+
+    def remove_completed_jobs(self):
+        for job in self.jobs[:]:  # Need to make a copy so we don't modify a list we're iterating over
+            job.remove_completed_pids()
+            if not job: self.jobs.remove(job)
+
+
+    @property
+    def process_count(self):
+        ''' return number of processes that currently ran by this driver instance
+        '''
+        return sum( map(len, self.jobs) )
+
+
+    def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+        return self.submit_serial_hpc_job(name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory, time, block, shell_wrapper)
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        cpu_usage = -time_module.time()/60./60.
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        def mfork():
+            ''' Check if number of child process is below cpu_count. And if it is - fork the new pocees and return its pid.
+            '''
+            while self.process_count >= self.cpu_count:
+                self.remove_completed_jobs()
+                if self.process_count >= self.cpu_count: time_module.sleep(.5)
+
+            sys.stdout.flush()
+            pid = os.fork()
+            # appending at caller level insted  if pid: self.jobs.append(pid) # We are parent!
+            return pid
+
+        current_job = self.JobID()
+        process = 0
+        for i in range(jobs_to_queue):
+
+            pid = mfork()
+            if not pid: # we are child process
+                command_line = 'cd {} && {} {}'.format(working_dir, executable, arguments.format(process=process) )
+                exit_code, log = execute('Running job {}.{}...'.format(name, i), command_line, tracer=self.tracer, return_='tuple')
+                with codecs.open(log_dir+'/.hpc.{name}.{i:02d}.log'.format(**vars()), 'w', encoding='utf-8', errors='replace') as f:
+                    f.write(command_line+'\n'+log)
+                    if exit_code:
+                        error_report = f'\n\n{command_line}\nERROR: PROCESS {name}.{i:02d} TERMINATED WITH NON-ZERO-EXIT-CODE {exit_code}!\n'
+                        f.write(error_report)
+                        print(log, error_report)
+
+                sys.exit(0)
+
+            else: # we are parent!
+                current_job.add_pid(pid)
+                # Need to potentially re-add to list, as remove_completed_jobs() might trim it.
+                if current_job not in self.jobs: self.jobs.append(current_job)
+
+            process += 1
+
+        if block:
+            #for p in all_queued_jobs: os.waitpid(p, 0)  # waiting for all child process to termintate...
+
+            self.wait_until_complete(current_job)
+            self.remove_completed_jobs()
+
+            cpu_usage += time_module.time()/60./60.
+            self.cpu_usage += cpu_usage * jobs_to_queue  # approximation...
+
+            current_job = self.JobID()
+
+        return current_job
+
+
+    @property
+    def number_of_cpu_per_node(self): return self.cpu_count
+
+
+    @property
+    def maximum_number_of_mpi_cpu(self): return self.cpu_count
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, memory=512, time=12, block=True, process_coefficient="1", requested_nodes=1, requested_processes_per_node=1):
+
+        if requested_nodes > 1:
+            print( "WARNING: " + str( requested_nodes ) + " nodes were requested, but we're running locally, so only 1 node will be used." )
+
+        if requested_processes_per_node > self.cpu_count:
+            print( "WARNING: " + str(requested_processes_per_node) + " processes were requested, but I only have " + str(self.cpu_count) + " CPUs.  Will launch " + str(self.cpu_count) + " processes."  )
+        actual_processes = min( requested_processes_per_node, self.cpu_count )
+
+        cpu_usage = -time_module.time()/60./60.
+
+        arguments = arguments.format(process=0)
+
+        command_line = f'cd {working_dir} && mpirun -np {actual_processes} {executable} {arguments}'
+        log = execute(f'Running job {name}...', command_line, tracer=self.tracer, return_='output')
+        with codecs.open(log_dir+'/.hpc.{name}.log'.format(**vars()), 'w', encoding='utf-8', errors='replace') as f: f.write(command_line+'\n'+log)
+
+        cpu_usage += time_module.time()/60./60.
+        self.cpu_usage += cpu_usage * actual_processes  # approximation...
+
+        # return None - we do not return anything from this version of submit which imply returning None which in turn will be treated as job-id for already finished job
+
+
+    def complete(self, job_id):
+        ''' Return job completion status. Return True if job completed and False otherwise
+        '''
+        self.remove_completed_jobs()
+        return job_id not in self.jobs
+
+
+    def cancel_job(self, job):
+        job.cancel();
+        if job in self.jobs:
+            self.jobs.remove(job)
+
+
+    def __repr__(self):
+        return 'MultiCore_HPC_Driver<>'

.rosetta-ci/hpc_drivers/slurm.py

@@ -0,0 +1,176 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import os, sys, time, collections, math
+import stat as stat_module
+
+
+try:
+    from .base import *
+
+except ImportError: # workaround for B2 back-end's
+    import imp
+    imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) + '/base.py')  # A bit of Python magic here, what we trying to say is this: from base import *, but path to base is calculated from our source location  # from base import HPC_Driver, execute, NT
+
+
+_T_slurm_array_job_template_ = '''\
+#!/bin/bash
+#
+#SBATCH --job-name={name}
+#SBATCH --output={log_dir}/.hpc.%x.%a.output
+#
+#SBATCH --time={time}:00
+#SBATCH --mem-per-cpu={memory}M
+#SBATCH --chdir={working_dir}
+#
+#SBATCH --array=1-{jobs_to_queue}
+
+srun {executable} {arguments}
+'''
+
+_T_slurm_mpi_job_template_ = '''\
+#!/bin/bash
+#
+#SBATCH --job-name={name}
+#SBATCH --output={log_dir}/.hpc.%x.output
+#
+#SBATCH --time={time}:00
+#SBATCH --mem-per-cpu={memory}M
+#SBATCH --chdir={working_dir}
+#
+#SBATCH --ntasks={ntasks}
+
+mpirun {executable} {arguments}
+'''
+
+class Slurm_HPC_Driver(HPC_Driver):
+    def head_node_execute(self, message, command_line, *args, **kwargs):
+        head_node = self.config['slurm'].get('head_node')
+
+        command_line, host = (f"ssh {head_node} cd `pwd` '&& {command_line}'", head_node) if head_node else (command_line, 'localhost')
+        return execute(f'Executiong on {host}: {message}' if message else '', command_line, *args, **kwargs)
+
+
+    # NodeGroup = collections.namedtuple('NodeGroup', 'nodes cores')
+
+    # @property
+    # def mpi_topology(self):
+    #     ''' return list of NodeGroup's
+    #     '''
+    #     pass
+
+
+    # @property
+    # def number_of_cpu_per_node(self): return int( self.config['condor']['mpi_cpu_per_node'] )
+
+    # @property
+    # def maximum_number_of_mpi_cpu(self):
+    #     return self.number_of_cpu_per_node * int( self.config['condor']['mpi_maximum_number_of_nodes'] )
+
+
+    # def complete(self, condor_job_id):
+    #     ''' Return job completion status. Note that single hpc_job may contatin inner list of individual HPC jobs, True should be return if they all run in to completion.
+    #     '''
+
+    #     execute('Releasing condor jobs...', 'condor_release $USER', return_='tuple')
+
+    #     s = execute('', 'condor_q $USER | grep $USER | grep {}'.format(condor_job_id), return_='output', terminate_on_failure=False).replace(' ', '').replace('\n', '')
+    #     if s: return False
+
+    #         # #setDaemonStatusAndPing('[Job #%s] Running... %s condor job(s) in queue...' % (self.id, len(s.split('\n') ) ) )
+    #         # n_jobs = len(s.split('\n'))
+    #         # s, o = execute('', 'condor_userprio -all | grep $USER@', return_='tuple')
+    #         # if s == 0:
+    #         #     jobs_running = o.split()
+    #         #     jobs_running = 'XX' if len(jobs_running) < 4 else jobs_running[4]
+    #         #     self.set_daemon_message("Waiting for condor to finish HPC jobs... [{} jobs in HPC-Queue, {} CPU's used]".format(n_jobs, jobs_running) )
+    #         #     print "{} condor jobs in queue... Sleeping 32s...    \r".format(n_jobs),
+    #         # sys.stdout.flush()
+    #         # time.sleep(32)
+    #     else:
+
+    #         #self.tracer('Waiting for condor to finish the jobs... DONE')
+    #         self.jobs.remove(condor_job_id)
+    #         self.cpu_usage += self.get_condor_accumulated_usage()
+    #         return True  # jobs already finished, we return empty list to prevent double counting of cpu_usage
+
+
+    def complete(self, slurm_job_id):
+        ''' Return True if job with given id is complete
+        '''
+
+        s = self.head_node_execute('', f'squeue -j {slurm_job_id} --noheader', return_='output', terminate_on_failure=False, silent=True)
+        if s: return False
+        else:
+            #self.tracer('Waiting for condor to finish the jobs... DONE')
+            self.jobs.remove(slurm_job_id)
+            return True  # jobs already finished, we return empty list to prevent double counting of cpu_usage
+
+
+    def cancel_job(self, slurm_job_id):
+        self.head_node_execute(f'Slurm_HPC_Driver.canceling job {slurm_job_id}...', f'scancel {slurm_job_id}', terminate_on_failure=False)
+
+
+    # def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+    #     print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+    #     return self.submit_serial_hpc_job(name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory, time, block, shell_wrapper)
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+
+        arguments = arguments.format(process='%a') # %a is SLURM array index
+        time = int( math.ceil(time*60) )
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        slurm_file = working_dir + f'/.hpc.{name}.slurm'
+
+        with open(slurm_file, 'w') as f: f.write( _T_slurm_array_job_template_.format( **vars() ) )
+
+
+        slurm_job_id = self.head_node_execute('Submitting SLURM array job...', f'cd {self.working_dir} && sbatch {slurm_file}',
+                                              tracer=self.tracer, return_='output'
+                                              ).split()[-1] # expecting something like `Submitted batch job 6122` in output
+
+
+        self.jobs.append(slurm_job_id)
+
+        if block:
+            self.wait_until_complete( [slurm_job_id] )
+            return None
+
+        else: return slurm_job_id
+
+
+
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, ntasks, memory=512, time=12, block=True, shell_wrapper=False):
+        ''' submit jobs as MPI job
+        '''
+        arguments = arguments.format(process='0')
+        time = int( math.ceil(time*60) )
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        slurm_file = working_dir + f'/.hpc.{name}.slurm'
+
+        with open(slurm_file, 'w') as f: f.write( _T_slurm_mpi_job_template_.format( **vars() ) )
+
+        slurm_job_id = self.head_node_execute('Submitting SLURM mpi job...', f'cd {self.working_dir} && sbatch {slurm_file}',
+                                              tracer=self.tracer, return_='output'
+                                              ).split()[-1] # expecting something like `Submitted batch job 6122` in output
+
+        self.jobs.append(slurm_job_id)
+
+        if block:
+            self.wait_until_complete( [slurm_job_id] )
+            return None
+
+        else: return slurm_job_id

.rosetta-ci/test-sets.yaml

@@ -0,0 +1,65 @@
+# map platform-string → platform definiton
+platforms:
+  ubuntu-20.04.gcc:
+    os: ubuntu-20.04
+    compiler: gcc
+    python: '3.9'
+
+  ubuntu-20.04.clang:
+    os: ubuntu-20.04
+    compiler: clang
+    python: '3.9'
+
+
+# map of test-set-name → tests
+test-sets:
+  main:
+    - ubuntu-20.04.clang.rfd
+
+  python:
+    - ubuntu-20.04.gcc.self.python
+    - ubuntu-20.04.clang.self.python
+
+  self:
+    - ubuntu-20.04.gcc.self.state
+    - ubuntu-20.04.gcc.self.subtests
+    - ubuntu-20.04.gcc.self.release
+
+
+
+# map of GitHub-label → [test-set]
+github-label-test-sets:
+  00 main: [main]
+  10 self: [self]
+  16 python: [python]
+
+
+# map of submit-page-category → tests
+# tests that does not get assigned will be automatically displayed in 'other' category
+category-tests:
+  main:
+    - rfd
+
+  self:
+    - self.state
+    - self.subtests
+    - self.release
+    - self.python
+
+
+# map branch → test-set to
+# specify list of tests that should be applied by-default during testing of each new commits to specific branch
+branch-test-sets:
+  main: [main]
+  benchmark: [main, python]
+
+
+# map branch → test-sets for pull-request's
+# specify which test-sets should be scheduled for PR's by-default (ie in addition to GH labels applied)
+# use empty branch name to specify defult value for (ie any branch not explicitly listed)
+pull-request-branch-test-sets:
+  # specific test sets for benchmark branch
+  benchmark: ['main', 'python']
+
+  # default, will apply to PR's to any other branch
+  '': ['main']

.rosetta-ci/tests/__init__.py

@@ -0,0 +1,765 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   tests/__init__.py
+## @brief  Common constats and types for all test types
+## @author Sergey Lyskov
+
+import os, time, sys, shutil, codecs, urllib.request, imp, subprocess, json, hashlib  # urllib.error, urllib.parse,
+import platform as  platform_module
+import types as types_module
+
+# ⚔ do not change wording below, it have to stay in sync with upstream (up to benchmark-model).
+# Copied from benchmark-model, standard state code's for tests results.
+
+__all__ = ['execute',
+           '_S_Values_', '_S_draft_', '_S_queued_', '_S_running_', '_S_passed_', '_S_failed_', '_S_build_failed_', '_S_script_failed_',
+           '_StateKey_', '_ResultsKey_', '_LogKey_', '_DescriptionKey_', '_TestsKey_',
+           '_multi_step_config_', '_multi_step_error_', '_multi_step_result_',
+           'to_bytes',
+]
+
+_S_draft_                 = 'draft'
+_S_queued_                = 'queued'
+_S_running_               = 'running'
+_S_passed_                = 'passed'
+_S_failed_                = 'failed'
+_S_build_failed_          = 'build failed'
+_S_script_failed_         = 'script failed'
+_S_queued_for_comparison_ = 'queued for comparison'
+
+_S_Values_ = [_S_draft_, _S_queued_, _S_running_, _S_passed_, _S_failed_, _S_build_failed_, _S_script_failed_, _S_queued_for_comparison_]
+
+_IgnoreKey_      = 'ignore'
+_StateKey_       = 'state'
+_ResultsKey_     = 'results'
+_LogKey_         = 'log'
+_DescriptionKey_ = 'description'
+_TestsKey_       = 'tests'
+_SummaryKey_     = 'summary'
+_FailedKey_      = 'failed'
+_TotalKey_       = 'total'
+_PlotsKey_       = 'plots'
+_FailedTestsKey_ = 'failed_tests'
+_HtmlKey_        = 'html'
+
+# file names for multi-step test files
+_multi_step_config_ = 'config.json'
+_multi_step_error_  = 'error.json'
+_multi_step_result_ = 'result.json'
+
+PyRosetta_unix_memory_requirement_per_cpu = 6  # Memory per sub-process in Gb's
+PyRosetta_unix_unit_test_memory_requirement_per_cpu = 3.0  # Memory per sub-process in Gb's for running PyRosetta unit tests
+
+# Commands to run all the scripts needed for setting up Rosetta compiles. (Run from main/source directory)
+PRE_COMPILE_SETUP_SCRIPTS = [ "./update_options.sh", "./update_submodules.sh", "./update_ResidueType_enum_files.sh", "python version.py" ]
+
+DEFAULT_PYTHON_VERSION='3.9'
+
+# Standard funtions and classes below ---------------------------------------------------------------------------------
+
+class BenchmarkError(Exception):
+    def __init__(self, value): self.value = value
+    def __repr__(self): return self.value
+    def __str__(self): return self.value
+
+
+class NT:  # named tuple
+    def __init__(self, **entries): self.__dict__.update(entries)
+    def __repr__(self):
+        r = 'NT: |'
+        for i in dir(self):
+            print(i)
+            if not i.startswith('__') and i != '_as_dict' and not isinstance(getattr(self, i), types_module.MethodType): r += '%s --> %s, ' % (i, getattr(self, i))
+        return r[:-2]+'|'
+
+    @property
+    def _as_dict(self):
+        return { a: getattr(self, a) for a in dir(self) if not a.startswith('__') and a != '_as_dict' and not isinstance(getattr(self, a), types_module.MethodType)}
+
+
+def Tracer(verbose=False):
+    return print if verbose else lambda x: None
+
+
+def to_unicode(b):
+    ''' Conver bytes to string and handle the errors. If argument is already in string - do nothing
+    '''
+    #return b if type(b) == unicode else unicode(b, 'utf-8', errors='replace')
+    return b if type(b) == str else str(b, 'utf-8', errors='backslashreplace')
+
+
+def to_bytes(u):
+    ''' Conver string to bytes and handle the errors. If argument is already of type bytes - do nothing
+    '''
+    return u if type(u) == bytes else u.encode('utf-8', errors='backslashreplace')
+
+
+''' Python-2 version
+def execute(message, commandline, return_=False, until_successes=False, terminate_on_failure=True, add_message_and_command_line_to_output=False):
+    message, commandline = to_unicode(message), to_unicode(commandline)
+
+    TR = Tracer()
+    TR(message);  TR(commandline)
+    while True:
+        (res, output) = commands.getstatusoutput(commandline)
+        # Subprocess results will always be a bytes-string.
+        # Probably ASCII, but may have some Unicode characters.
+        # A UTF-8 decode will probably get decent results 99% of the time
+        # and the replace option will gracefully handle the rest.
+        output = to_unicode(output)
+
+        TR(output)
+
+        if res and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        print( "Error while executing %s: %s\n" % (message, output) )
+        print( "Sleeping 60s... then I will retry..." )
+        time.sleep(60)
+
+    if add_message_and_command_line_to_output: output = message + '\nCommand line: ' + commandline + '\n' + output
+
+    if return_ == 'tuple': return(res, output)
+
+    if res and terminate_on_failure:
+        TR("\nEncounter error while executing: " + commandline)
+        if return_==True: return res
+        else:
+            print("\nEncounter error while executing: " + commandline + '\n' + output)
+            raise BenchmarkError("\nEncounter error while executing: " + commandline + '\n' + output)
+
+    if return_ == 'output': return output
+    else: return res
+'''
+
+def execute_through_subprocess(command_line):
+    # exit_code, output = subprocess.getstatusoutput(command_line)
+
+    # p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    # output, errors = p.communicate()
+    # output = (output + errors).decode(encoding='utf-8', errors='backslashreplace')
+    # exit_code = p.returncode
+
+    # previous 'main' version based on subprocess module. Main issue that output of segfaults will not be captured since they generated by shell
+    p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    output, errors = p.communicate()
+    # output = output + errors # ← we redirected stderr into same pipe as stdcout so errors is None, - no need to concatenate
+    output = output.decode(encoding='utf-8', errors='backslashreplace')
+    exit_code = p.returncode
+
+    return exit_code, output
+
+
+def execute_through_pexpect(command_line):
+    import pexpect
+
+    child = pexpect.spawn('/bin/bash', ['-c', command_line])
+    child.expect(pexpect.EOF)
+    output = child.before.decode(encoding='utf-8', errors='backslashreplace')
+    child.close()
+    exit_code = child.signalstatus or child.exitstatus
+
+    return exit_code, output
+
+
+def execute_through_pty(command_line):
+    import pty, select
+
+    if sys.platform == "darwin":
+
+        master, slave = pty.openpty()
+        p = subprocess.Popen(command_line, shell=True, stdout=slave, stdin=slave,
+                             stderr=subprocess.STDOUT, close_fds=True)
+
+        buffer = []
+        while True:
+            try:
+                if select.select([master], [], [], 0.2)[0]:  # has something to read
+                    data = os.read(master, 1 << 22)
+                    if data: buffer.append(data)
+
+                elif (p.poll() is not None)  and  (not select.select([master], [], [], 0.2)[0] ): break  # process is finished and output buffer if fully read
+
+            except OSError: break  # OSError will be raised when child process close PTY descriptior
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+
+        os.close(master)
+        os.close(slave)
+
+        p.wait()
+        exit_code = p.returncode
+
+        '''
+        buffer = []
+        while True:
+            if select.select([master], [], [], 0.2)[0]:  # has something to read
+                data = os.read(master, 1 << 22)
+                if data: buffer.append(data)
+                # else: break  # # EOF - well, technically process _should_ be finished here...
+
+            # elif time.sleep(1) or (p.poll() is not None): # process is finished (sleep here is intentional to trigger race condition, see solution for this on the next few lines)
+            #     assert not select.select([master], [], [], 0.2)[0]  # should be nothing left to read...
+            #     break
+
+            elif (p.poll() is not None)  and  (not select.select([master], [], [], 0.2)[0] ): break  # process is finished and output buffer if fully read
+
+        assert not select.select([master], [], [], 0.2)[0]  # should be nothing left to read...
+
+        os.close(slave)
+        os.close(master)
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+        exit_code = p.returncode
+        '''
+
+    else:
+
+        master, slave = pty.openpty()
+        p = subprocess.Popen(command_line, shell=True, stdout=slave, stdin=slave,
+                             stderr=subprocess.STDOUT, close_fds=True)
+
+        os.close(slave)
+
+        buffer = []
+        while True:
+            try:
+                data = os.read(master, 1 << 22)
+                if data: buffer.append(data)
+            except OSError: break  # OSError will be raised when child process close PTY descriptior
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+
+        os.close(master)
+
+        p.wait()
+        exit_code = p.returncode
+
+    return exit_code, output
+
+
+
+def execute(message, command_line, return_='status', until_successes=False, terminate_on_failure=True, silent=False, silence_output=False, silence_output_on_errors=False, add_message_and_command_line_to_output=False):
+    if not silent: print(message);  print(command_line); sys.stdout.flush();
+    while True:
+
+        #exit_code, output = execute_through_subprocess(command_line)
+        #exit_code, output = execute_through_pexpect(command_line)
+        exit_code, output = execute_through_pty(command_line)
+
+        if (exit_code  and  not silence_output_on_errors) or  not (silent or silence_output): print(output); sys.stdout.flush();
+
+        if exit_code and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        print( "Error while executing {}: {}\n".format(message, output) )
+        print("Sleeping 60s... then I will retry...")
+        sys.stdout.flush();
+        time.sleep(60)
+
+    if add_message_and_command_line_to_output: output = message + '\nCommand line: ' + command_line + '\n' + output
+
+    if return_ == 'tuple'  or  return_ == tuple: return(exit_code, output)
+
+    if exit_code and terminate_on_failure:
+        print("\nEncounter error while executing: " + command_line)
+        if return_==True: return True
+        else:
+            print('\nEncounter error while executing: ' + command_line + '\n' + output);
+            raise BenchmarkError('\nEncounter error while executing: ' + command_line + '\n' + output)
+
+    if return_ == 'output': return output
+    else: return exit_code
+
+
+def parallel_execute(name, jobs, rosetta_dir, working_dir, cpu_count, time=16):
+    ''' Execute command line in parallel on local host
+        time specifies the upper limit for cpu-usage runtime (in minutes) for any one process in the parallel execution.
+
+        jobs should be dict with following structure:
+        {
+            'job-string-id-1’: command_line-1,
+            'job-string-id-2’: command_line-2,
+            ...
+        }
+
+        return: dict with jobs-id's as keys and value as dict with 'output' and 'result' keys:
+        {
+            "job-string-id-1": {
+                "output": "stdout + stdderr output of command_line-1",
+                "result": <integer exit code for command_line-1>
+            },
+            "c2": {
+                "output": "stdout + stdderr output of command_line-2",
+                "result": <integer exit code for command_line-2>
+            },
+            ...
+        }
+    '''
+    job_file_name = working_dir + '/' + name
+    with open(job_file_name + '.json', 'w') as f: json.dump(jobs, f, sort_keys=True, indent=2) # JSON handles unicode internally
+    if time is not None:
+        allowed_time = int(time*60)
+        ulimit_command = f'ulimit -t {allowed_time} && '
+    else:
+        ulimit_command = ''
+    command = f'cd {working_dir} && ' + ulimit_command + f'{rosetta_dir}/tests/benchmark/util/parallel.py -j{cpu_count} {job_file_name}.json'
+    execute("Running {} in parallel with {} CPU's...".format(name, cpu_count), command )
+
+    with open(job_file_name+'.results.json') as f: return json.load(f)
+
+
+def calculate_unique_prefix_path(platform, config):
+    ''' calculate path for prefix location that is unique for this machine and OS
+    '''
+    hostname = os.uname()[1]
+    return config['prefix'] + '/' + hostname + '/' + platform['os']
+
+
+def get_python_include_and_lib(python):
+    ''' calculate python include dir and lib dir from given python executable path
+    '''
+    #python = os.path.realpath(python)
+    python_bin_dir = python.rpartition('/')[0]
+    python_config = f'{python} {python}-config' if python.endswith('2.7') else f'{python}-config'
+
+    #if not os.path.isfile(python_config): python_config = python_bin_dir + '/python-config'
+
+    info = execute('Getting python configuration info...', f'unset __PYVENV_LAUNCHER__ && cd {python_bin_dir} && PATH=.:$PATH && {python_config} --prefix --includes', return_='output').replace('\r', '').split('\n')  # Python-3 only: --abiflags
+    python_prefix = info[0]
+    python_include_dir = info[1].split()[0][len('-I'):]
+    python_lib_dir = python_prefix + '/lib'
+    #python_abi_suffix = info[2]
+    #print(python_include_dir, python_lib_dir)
+
+    return NT(python_include_dir=python_include_dir, python_lib_dir=python_lib_dir)
+
+
+def local_open_ssl_install(prefix, build_prefix, jobs):
+    ''' install OpenSSL at given prefix, return url of source archive
+    '''
+    #with tempfile.TemporaryDirectory('open_ssl_build', dir=prefix) as build_prefix:
+
+    url = 'https://www.openssl.org/source/openssl-1.1.1b.tar.gz'
+    #url = 'https://www.openssl.org/source/openssl-3.0.0.tar.gz'
+
+
+    archive = build_prefix + '/' + url.split('/')[-1]
+    build_dir = archive.rpartition('.tar.gz')[0]
+    if os.path.isdir(build_dir): shutil.rmtree(build_dir)
+
+    with open(archive, 'wb') as f:
+        response = urllib.request.urlopen(url)
+        f.write( response.read() )
+
+    execute('Unpacking {}'.format(archive), 'cd {build_prefix} && tar -xvzf {archive}'.format(**vars()) )
+
+    execute('Configuring...', f'cd {build_dir} && ./config --prefix={prefix}')
+    execute('Building...',    f'cd {build_dir} && make -j{jobs}')
+    execute('Installing...',  f'cd {build_dir} && make -j{jobs} install')
+
+    return url
+
+
+def remove_pip_and_easy_install(prefix_root_path):
+    ''' remove `pip` and `easy_install` executable from given Python / virtual-environments install
+    '''
+    for f in os.listdir(prefix_root_path + '/bin'):  # removing all pip's and easy_install's to make sure that environment is immutable
+        for p in ['pip', 'easy_install']:
+            if f.startswith(p): os.remove(prefix_root_path + '/bin/' + f)
+
+
+
+def local_python_install(platform, config):
+    ''' Perform local install of given Python version and return path-to-python-interpreter, python_include_dir, python_lib_dir
+        If previous install is detected skip installiation.
+        Provided Python install will _persistent_ and _immutable_
+    '''
+    jobs = config['cpu_count']
+    compiler, cpp_compiler = ('clang', 'clang++') if platform['os'] == 'mac' else ('gcc', 'g++')  # disregarding platform compiler setting and instead use default compiler for platform
+
+    python_version = platform.get('python', DEFAULT_PYTHON_VERSION)
+
+    if python_version.endswith('.s'):
+        assert python_version == f'{sys.version_info.major}.{sys.version_info.minor}.s'
+        #root = executable.rpartition('/bin/python')[0]
+        h = hashlib.md5(); h.update( (sys.executable + sys.version).encode('utf-8', errors='backslashreplace') ); hash = h.hexdigest()
+        return NT(
+            python = sys.executable,
+            root = None,
+            python_include_dir = None,
+            python_lib_dir = None,
+            version = python_version,
+            url = None,
+            platform = platform,
+            config = config,
+            hash = hash,
+        )
+
+    # deprecated, no longer needed
+    # python_version = {'python2'   : '2.7',
+    #                   'python2.7' : '2.7',
+    #                   'python3'   : '3.5',
+    # }.get(python_version, python_version)
+
+    # for security reasons we only allow installs for version listed here with hand-coded URL's
+    python_sources = {
+        '2.7' : 'https://www.python.org/ftp/python/2.7.18/Python-2.7.18.tgz',
+
+        '3.5'  : 'https://www.python.org/ftp/python/3.5.9/Python-3.5.9.tgz',
+        '3.6'  : 'https://www.python.org/ftp/python/3.6.15/Python-3.6.15.tgz',
+        '3.7'  : 'https://www.python.org/ftp/python/3.7.14/Python-3.7.14.tgz',
+        '3.8'  : 'https://www.python.org/ftp/python/3.8.14/Python-3.8.14.tgz',
+        '3.9'  : 'https://www.python.org/ftp/python/3.9.14/Python-3.9.14.tgz',
+        '3.10' : 'https://www.python.org/ftp/python/3.10.10/Python-3.10.10.tgz',
+        '3.11' : 'https://www.python.org/ftp/python/3.11.2/Python-3.11.2.tgz',
+    }
+
+    # map of env -> ('shell-code-before ./configure', 'extra-arguments-for-configure')
+    extras = {
+        #('mac',) :          ('__PYVENV_LAUNCHER__="" MACOSX_DEPLOYMENT_TARGET={}'.format(platform_module.mac_ver()[0]), ''),  # __PYVENV_LAUNCHER__ now used by-default for all platform installs
+        ('mac',) :          ('MACOSX_DEPLOYMENT_TARGET={}'.format(platform_module.mac_ver()[0]), ''),
+        ('linux',  '2.7') : ('', '--enable-unicode=ucs4'),
+        ('ubuntu', '2.7') : ('', '--enable-unicode=ucs4'),
+    }
+
+    #packages = '' if (python_version[0] == '2' or  python_version == '3.5' ) and  platform['os'] == 'mac' else 'pip setuptools wheel' # 2.7 is now deprecated on Mac so some packages could not be installed
+    packages = 'setuptools'
+
+    url = python_sources[python_version]
+
+    extra = extras.get( (platform['os'],)  , ('', '') )
+    extra = extras.get( (platform['os'], python_version) , extra)
+
+    extra = ('unset __PYVENV_LAUNCHER__ && ' + extra[0], extra[1])
+
+    options = '--with-ensurepip' #'--without-ensurepip'
+    signature = f'v1.5.1 url: {url}\noptions: {options}\ncompiler: {compiler}\nextra: {extra}\npackages: {packages}\n'
+
+    h = hashlib.md5(); h.update( signature.encode('utf-8', errors='backslashreplace') ); hash = h.hexdigest()
+
+    root = calculate_unique_prefix_path(platform, config) + '/python-' + python_version + '.' +  compiler + '/' + hash
+
+    signature_file_name = root + '/.signature'
+
+    #activate   = root + '/bin/activate'
+    executable = root + '/bin/python' + python_version
+
+    # if os.path.isfile(executable)  and  (not execute('Getting python configuration info...', '{executable}-config --prefix --includes'.format(**vars()), terminate_on_failure=False) ):
+    #     print('found executable!')
+    #     _, executable_version = execute('Checking Python interpreter version...', '{executable} --version'.format(**vars()), return_='tuple')
+    #     executable_version = executable_version.split()[-1]
+    # else: executable_version = ''
+    # print('executable_version: {}'.format(executable_version))
+    #if executable_version != url.rpartition('Python-')[2][:-len('.tgz')]:
+
+    if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature:
+        #print('Install for Python-{} is detected, skipping installation procedure...'.format(python_version))
+        pass
+
+    else:
+        print( 'Installing Python-{python_version}, using {url} with extra:{extra}...'.format( **vars() ) )
+
+        if os.path.isdir(root): shutil.rmtree(root)
+
+        build_prefix = os.path.abspath(root + '/../build-python-{}'.format(python_version) )
+
+        if not os.path.isdir(root): os.makedirs(root)
+        if not os.path.isdir(build_prefix): os.makedirs(build_prefix)
+
+        platform_is_mac = True if platform['os'] in ['mac', 'm1'] else False
+        platform_is_linux = not platform_is_mac
+
+        #if False and platform['os'] == 'mac' and platform_module.machine() == 'arm64' and tuple( map(int, python_version.split('.') ) ) >= (3, 9):
+        if ( platform['os'] == 'mac' and python_version == '3.6' ) \
+           or ( platform_is_linux and python_version in ['3.10', '3.11'] ):
+            open_ssl_url = local_open_ssl_install(root, build_prefix, jobs)
+            options += f' --with-openssl={root} --with-openssl-rpath=auto'
+            #signature += 'OpenSSL install: ' + open_ssl_url + '\n'
+
+        archive = build_prefix + '/' + url.split('/')[-1]
+        build_dir = archive.rpartition('.tgz')[0]
+        if os.path.isdir(build_dir): shutil.rmtree(build_dir)
+
+        with open(archive, 'wb') as f:
+            #response = urllib2.urlopen(url)
+            response = urllib.request.urlopen(url)
+            f.write( response.read() )
+
+        #execute('Execution environment:', 'env'.format(**vars()) )
+
+        execute('Unpacking {}'.format(archive), 'cd {build_prefix} && tar -xvzf {archive}'.format(**vars()) )
+
+        #execute('Building and installing...', 'cd {} && CC={compiler} CXX={cpp_compiler} {extra[0]} ./configure {extra[1]} --prefix={root} && {extra[0]} make -j{jobs} && {extra[0]} make install'.format(build_dir, **locals()) )
+        execute('Configuring...', 'cd {} && CC={compiler} CXX={cpp_compiler} {extra[0]} ./configure {options} {extra[1]} --prefix={root}'.format(build_dir, **locals()) )
+        execute('Building...', 'cd {} && {extra[0]} make -j{jobs}'.format(build_dir, **locals()) )
+        execute('Installing...', 'cd {} && {extra[0]} make -j{jobs} install'.format(build_dir, **locals()) )
+
+        shutil.rmtree(build_prefix)
+
+        #execute('Updating setuptools...', f'cd {root} && {root}/bin/pip{python_version} install --upgrade setuptools wheel' )
+
+        # if 'certifi' not in packages:
+        #     packages += ' certifi'
+
+        if packages: execute( f'Installing packages {packages}...', f'cd {root} && unset __PYVENV_LAUNCHER__ && {root}/bin/pip{python_version} install --upgrade {packages}' )
+        #if packages: execute( f'Installing packages {packages}...', f'cd {root} && unset __PYVENV_LAUNCHER__ && {executable} -m pip install --upgrade {packages}' )
+
+        remove_pip_and_easy_install(root)  # removing all pip's and easy_install's to make sure that environment is immutable
+
+        with open(signature_file_name, 'w') as f: f.write(signature)
+
+        print( 'Installing Python-{python_version}, using {url} with extra:{extra}... Done.'.format( **vars() ) )
+
+    il = get_python_include_and_lib(executable)
+
+    return NT(
+        python = executable,
+        root = root,
+        python_include_dir = il.python_include_dir,
+        python_lib_dir = il.python_lib_dir,
+        version = python_version,
+        url = url,
+        platform = platform,
+        config = config,
+        hash = hash,
+    )
+
+
+
+def setup_python_virtual_environment(working_dir, python_environment, packages=''):
+    ''' Deploy Python virtual environment at working_dir
+    '''
+
+    python = python_environment.python
+
+    execute('Setting up Python virtual environment...', 'unset __PYVENV_LAUNCHER__ && {python} -m venv --clear {working_dir}'.format(**vars()) )
+
+    activate = f'unset __PYVENV_LAUNCHER__ && . {working_dir}/bin/activate'
+
+    bin=working_dir+'/bin'
+
+    if packages: execute('Installing packages: {}...'.format(packages), 'unset __PYVENV_LAUNCHER__ && {bin}/python {bin}/pip install --upgrade pip setuptools && {bin}/python {bin}/pip install --progress-bar off {packages}'.format(**vars()) )
+    #if packages: execute('Installing packages: {}...'.format(packages), '{bin}/pip{python_environment.version} install {packages}'.format(**vars()) )
+
+    return NT(activate = activate, python = bin + '/python', root = working_dir, bin = bin)
+
+
+
+def setup_persistent_python_virtual_environment(python_environment, packages):
+    ''' Setup _persistent_ and _immutable_ Python virtual environment which will be saved between test runs
+    '''
+
+    if python_environment.version.startswith('2.'):
+        assert not packages, f'ERROR: setup_persistent_python_virtual_environment does not support Python-2.* with non-empty package list!'
+        return NT(activate = ':', python = python_environment.python, root = python_environment.root, bin = python_environment.root + '/bin')
+
+    else:
+        #if 'certifi' not in packages: packages += ' certifi'
+
+        h = hashlib.md5()
+        h.update(f'v1.0.0 platform: {python_environment.platform} python_source_url: {python_environment.url} python-hash: {python_environment.hash} packages: {packages}'.encode('utf-8', errors='backslashreplace') )
+        hash = h.hexdigest()
+
+        prefix = calculate_unique_prefix_path(python_environment.platform, python_environment.config)
+
+        root = os.path.abspath( prefix + '/python_virtual_environments/' + '/python-' + python_environment.version + '/' + hash )
+        signature_file_name = root + '/.signature'
+        signature = f'setup_persistent_python_virtual_environment v1.0.0\npython: {python_environment.hash}\npackages: {packages}\n'
+
+        activate = f'unset __PYVENV_LAUNCHER__ && . {root}/bin/activate'
+        bin = f'{root}/bin'
+
+        if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature: pass
+        else:
+            if os.path.isdir(root): shutil.rmtree(root)
+            setup_python_virtual_environment(root, python_environment, packages=packages)
+            remove_pip_and_easy_install(root)  # removing all pip's and easy_install's to make sure that environment is immutable
+            with open(signature_file_name, 'w') as f: f.write(signature)
+
+        return NT(activate = activate, python = bin + '/python', root = root, bin = bin, hash = hash)
+
+
+
+def _get_path_to_conda_root(platform, config):
+    ''' Perform local (prefix) install of miniconda and return NT(activate, conda_root_dir, conda)
+        this function is for inner use only, - to setup custom conda environment inside your test use `setup_conda_virtual_environment` defined below
+    '''
+    miniconda_sources = {
+        'mac'    : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh',
+        'linux'  : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh',
+        'aarch64': 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh',
+        'ubuntu' : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh',
+        'm1'     : 'https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.1-MacOSX-arm64.sh',
+    }
+
+    conda_sources = {
+        'mac'    : 'https://repo.continuum.io/archive/Anaconda3-2018.12-MacOSX-x86_64.sh',
+        'linux'  : 'https://repo.continuum.io/archive/Anaconda3-2018.12-Linux-x86_64.sh',
+        'ubuntu' : 'https://repo.continuum.io/archive/Anaconda3-2018.12-Linux-x86_64.sh',
+    }
+
+    #platform_os = 'm1' if platform_module.machine() == 'arm64' else platform['os']
+    #url = miniconda_sources[ platform_os ]
+
+    platform_os = platform['os']
+    for o in 'alpine centos ubuntu'.split():
+        if platform_os.startswith(o): platform_os = 'linux'
+
+    url = miniconda_sources[platform_os]
+
+    version = '1'
+    channels = ''  # conda-forge
+
+    #packages = ['conda-build gcc libgcc', 'libgcc=5.2.0'] # libgcc installs is workaround for "Anaconda libstdc++.so.6: version `GLIBCXX_3.4.20' not found", see: https://stackoverflow.com/questions/48453497/anaconda-libstdc-so-6-version-glibcxx-3-4-20-not-found
+    #packages = ['conda-build gcc'] # libgcc installs is workaround for "Anaconda libstdc++.so.6: version `GLIBCXX_3.4.20' not found", see: https://stackoverflow.com/questions/48453497/anaconda-libstdc-so-6-version-glibcxx-3-4-20-not-found
+    packages = ['conda-build anaconda-client conda-verify',]
+
+    signature = f'url: {url}\nversion: {version}\channels: {channels}\npackages: {packages}\n'
+
+    root = calculate_unique_prefix_path(platform, config) + '/conda'
+
+    signature_file_name = root + '/.signature'
+
+    # presense of __PYVENV_LAUNCHER__,PYTHONHOME, PYTHONPATH sometimes confuse Python so we have to unset them
+    unset = 'unset __PYVENV_LAUNCHER__ && unset PYTHONHOME && unset PYTHONPATH'
+    activate = unset + ' && . ' + root + '/bin/activate'
+
+    executable = root + '/bin/conda'
+
+
+    if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature:
+        print( f'Install for MiniConda is detected, skipping installation procedure...' )
+
+    else:
+        print( f'Installing MiniConda, using {url}...' )
+
+        if os.path.isdir(root): shutil.rmtree(root)
+
+        build_prefix = os.path.abspath(root + f'/../build-conda' )
+
+        #if not os.path.isdir(root): os.makedirs(root)
+        if not os.path.isdir(build_prefix): os.makedirs(build_prefix)
+
+        archive = build_prefix + '/' + url.split('/')[-1]
+
+        with open(archive, 'wb') as f:
+            response = urllib.request.urlopen(url)
+            f.write( response.read() )
+
+        execute('Installing conda...', f'cd {build_prefix} && {unset} && bash {archive} -b -p {root}' )
+
+        # conda update --yes --quiet -n base -c defaults conda
+
+        if channels: execute(f'Adding extra channles {channels}...', f'cd {build_prefix} && {activate} && conda config --add channels {channels}' )
+
+        for p in packages: execute(f'Installing conda packages: {p}...', f'cd {build_prefix} && {activate} && conda install --quiet --yes {p}' )
+
+        shutil.rmtree(build_prefix)
+
+        with open(signature_file_name, 'w') as f: f.write(signature)
+
+        print( f'Installing MiniConda, using {url}... Done.' )
+
+    execute(f'Updating conda base...', f'{activate} && conda update --all --yes' )
+    return NT(conda=executable, root=root, activate=activate, url=url)
+
+
+
+def setup_conda_virtual_environment(working_dir, platform, config, packages=''):
+    ''' Deploy Conda virtual environment at working_dir
+    '''
+    conda_root_env = _get_path_to_conda_root(platform, config)
+    activate = conda_root_env.activate
+
+    python_version = platform.get('python', DEFAULT_PYTHON_VERSION)
+
+    prefix = os.path.abspath( working_dir + '/.conda-python-' + python_version )
+
+    command_line = f'conda create --quiet --yes --prefix {prefix} python={python_version}'
+
+    execute( f'Setting up Conda for Python-{python_version} virtual environment...', f'cd {working_dir} && {activate} && ( {command_line} || ( conda clean --yes && {command_line} ) )' )
+
+    activate = f'{activate} && conda activate {prefix}'
+
+    if packages: execute( f'Setting up extra packages {packages}...', f'cd {working_dir} && {activate} && conda install --quiet --yes {packages}' )
+
+    python = prefix + '/bin/python' + python_version
+
+    il = get_python_include_and_lib(python)
+
+    return NT(
+        activate = activate,
+        root = prefix,
+        python = python,
+        python_include_dir = il.python_include_dir,
+        python_lib_dir = il.python_lib_dir,
+        version = python_version,
+        activate_base = conda_root_env.activate,
+        url = prefix, # conda_root_env.url,
+        platform=platform,
+        config=config,
+    )
+
+
+
+class FileLock():
+    ''' Implementation of file-lock object that could be use with Python `with` statement
+    '''
+
+    def __init__(self, file_name):
+        self.locked = False
+        self.file_name = file_name
+
+
+    def __enter__(self):
+        if not self.locked: self.acquire()
+        return self
+
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.locked: self.release()
+
+
+    def __del__(self):
+        self.release()
+
+
+    def acquire(self):
+        while True:
+            try:
+                os.close( os.open(self.file_name, os.O_CREAT | os.O_EXCL, mode=0o600) )
+                self.locked = True
+                break
+
+            except FileExistsError as e:
+                time.sleep(60)
+
+
+    def release(self):
+        if self.locked:
+            os.remove(self.file_name)
+            self.locked = False
+
+
+
+def convert_submodule_urls_from_ssh_to_https(repository_root):
+    ''' switching submodules URL to HTTPS so we can clone without SSH key
+    '''
+    with open(f'{repository_root}/.gitmodules') as f: m = f.read()
+    with open(f'{repository_root}/.gitmodules', 'w') as f:
+        f.write(
+            m
+            .replace('url = git@github.com:', 'url = https://github.com/')
+            .replace('url = ../../../',       'url = https://github.com/RosettaCommons/')
+            .replace('url = ../../',          'url = https://github.com/RosettaCommons/')
+            .replace('url = ../',             'url = https://github.com/RosettaCommons/')
+        )

.rosetta-ci/tests/rfd.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   rfd.py
+## @brief  main test files for RFdiffusion
+## @author Sergey Lyskov
+
+
+import imp
+imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) +  '/__init__.py')  # A bit of Python magic here, what we trying to say is this: from __init__ import *, but init is calculated from file location
+
+_api_version_ = '1.0'
+
+import os, tempfile, shutil
+import urllib.request
+
+
+_models_urls_ = '''
+http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt
+'''.split()
+
+
+def run_main_test_suite(repository_root, working_dir, platform, config, debug):
+    full_log = ''
+
+    python_environment = local_python_install(platform, config)
+
+    models_dir = repository_root + '/models'
+    if not os.path.isdir(models_dir): os.makedirs(models_dir)
+
+    for url in _models_urls_:
+        file_name = models_dir + '/' + url.split('/')[-1]
+        tmp_file_name = file_name + '.tmp'
+        if not os.path.isfile(file_name):
+            print(f'downloading {url}...')
+            full_log += f'downloading {url}...\n'
+            urllib.request.urlretrieve(url, tmp_file_name)
+            os.rename(tmp_file_name, file_name)
+
+    execute('unpacking ppi scaffolds...', f'cd {repository_root} && tar -xvf examples/ppi_scaffolds_subset.tar.gz -C examples')
+
+    with tempfile.TemporaryDirectory(dir=working_dir) as tmpdirname:
+    # tmpdirname = working_dir+'/.ve'
+    # if True:
+
+        #ve = setup_persistent_python_virtual_environment(python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl')
+        #ve = setup_python_virtual_environment(working_dir+'/.ve', python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl e3nn icecream pyrsistent wandb pynvml decorator jedi hydra-core')
+        ve = setup_python_virtual_environment(tmpdirname, python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl e3nn icecream pyrsistent wandb pynvml decorator jedi hydra-core')
+
+        execute('Installing local se3-transformer package...', f'cd {repository_root}/env/SE3Transformer && {ve.bin}/pip3 install --editable .')
+        execute('Installing RFdiffusion package...', f'cd {repository_root} && {ve.bin}/pip3 install --editable .')
+
+        #res, output = execute('running unit tests...', f'{ve.activate} && cd {repository_root} && python -m unittest', return_='tuple', add_message_and_command_line_to_output=True)
+        #res, output = execute('running unit tests...', f'cd {repository_root} && {ve.bin}/pytest', return_='tuple')
+
+
+        results_file = f'{repository_root}/tests/.results.json'
+        if os.path.isfile(results_file): os.remove(results_file)
+
+        res, output = execute('running RFdiffusion tests...', f'{ve.activate} && cd {repository_root}/tests && python test_diffusion.py', return_='tuple', add_message_and_command_line_to_output=True)
+
+        if os.path.isfile(results_file):
+            with open(results_file) as f: sub_tests_reults = json.load(f)
+
+            state = _S_passed_
+            for r in sub_tests_reults.values():
+                if r[_StateKey_] == _S_failed_:
+                    state = _S_failed_
+                    break
+
+        else:
+            sub_tests_reults = {}
+            output  += '\n\nEmpty sub-test results, marking test as `failed`...'
+            state = _S_failed_
+
+        shutil.move(f'{repository_root}/tests/outputs', f'{working_dir}/outputs')
+
+        for d in os.listdir(f'{repository_root}/tests'):
+            p = f'{repository_root}/tests/{d}'
+            if d.startswith('tests_') and os.path.isdir(p): shutil.rmtree(p)
+
+        results = {
+            _StateKey_ : state,
+            _LogKey_ : full_log + '\n' + output,
+            _ResultsKey_ : {
+                _TestsKey_ : sub_tests_reults,
+            },
+        }
+
+    return results
+
+
+
+def run(test, repository_root, working_dir, platform, config, hpc_driver=None, verbose=False, debug=False):
+    if test == '': return run_main_test_suite(repository_root=repository_root, working_dir=working_dir, platform=platform, config=config, debug=debug)
+    else: raise BenchmarkError('Unknow scripts test: {}!'.format(test))

.rosetta-ci/tests/self.md

@@ -0,0 +1,6 @@
+# self test suite
+These tests are design to help debug interface between testing server and Rosetta testing scripts
+
+-----
+### python
+Test Python platform support and functionality of local and persistent Python virtual environments

.rosetta-ci/tests/self.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   dummy.py
+## @brief  self-test and debug-aids tests
+## @author Sergey Lyskov
+
+import os, os.path, shutil, re, string
+import json
+
+import random
+
+import imp
+imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) +  '/__init__.py')  # A bit of Python magic here, what we trying to say is this: from __init__ import *, but init is calculated from file location
+
+_api_version_ = '1.0'
+
+
+def run_state_test(repository_root, working_dir, platform, config):
+    revision_id = config['revision']
+    states = (_S_passed_, _S_failed_, _S_build_failed_, _S_script_failed_)
+    state = states[revision_id % len(states)]
+
+    return {_StateKey_ : state,  _ResultsKey_ : {},  _LogKey_ : f'run_state_test: setting test state to {state!r}...' }
+
+
+sub_test_description_template = '''\
+# subtests_test test suite
+These sub-test description is generated for 3/4 of sub-tests
+
+-----
+### {name}
+The warm time, had already disappeared like dust. Broken rain, fragment of light shadow, bring more pain to my heart...
+-----
+'''
+
+def run_subtests_test(repository_root, working_dir, platform, config):
+    tests = {}
+    for i in range(16):
+        name = f's-{i:02}'
+        log = ('x'*63 + '\n') * 16 * 256 * i
+        s = i % 3
+        if   s == 0: state = _S_passed_
+        elif s == 1: state = _S_failed_
+        else:        state = _S_script_failed_
+
+        if i % 4:
+            os.mkdir( f'{working_dir}/{name}' )
+            with open(f'{working_dir}/{name}/description.md', 'w') as f: f.write( sub_test_description_template.format(**vars()) )
+
+            with open( f'{working_dir}/{name}/fantome.txt', 'w') as f: f.write('No one wants to hear the sequel to a fairytale\n')
+
+        tests[name] = { _StateKey_ : state, _LogKey_   : log, }
+
+    test_log = ('*'*63 + '\n') * 16 * 1024 * 16
+    return {_StateKey_ : _S_failed_,  _ResultsKey_ : {_TestsKey_: tests},  _LogKey_ : test_log }
+
+
+def run_regression_test(repository_root, working_dir, platform, config):
+    const     = 'const'
+    volatile  = 'volatile'
+    new       = ''.join( random.sample( string.ascii_letters + string.digits, 8) )
+    oversized = 'oversized'
+
+    sub_tests = [const, volatile, new]
+
+    const_dir = working_dir + '/' + const
+    os.mkdir(const_dir)
+    with open(const_dir + '/const_data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(32) ) ) )
+
+    volatile_dir = working_dir + '/' + volatile
+    os.mkdir(volatile_dir)
+    with open(volatile_dir + '/const_data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(32, 64) ) ) )
+    with open(volatile_dir + '/volatile_data', 'w') as f: f.write( '\n'.join( ( ''.join(random.sample( string.ascii_letters + string.digits, 8) ) for i in range(32) ) ) )
+
+    new_dir = working_dir + '/' + new
+    os.mkdir(new_dir)
+    with open(new_dir + '/data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(64)) ) )
+
+
+    new_dir = working_dir + '/' + oversized
+    os.mkdir(new_dir)
+    with open(new_dir + '/large', 'w') as f: f.write( ('x'*63 + '\n')*16*1024*256 +'extra')
+
+    return {_StateKey_ : _S_queued_for_comparison_,  _ResultsKey_ : {},  _LogKey_ : f'sub-tests: {sub_tests!r}' }
+
+
+
+def run_release_test(repository_root, working_dir, platform, config):
+    release_root = config['mounts'].get('release_root')
+
+    branch = config['branch']
+    revision = config['revision']
+
+    assert release_root, "config['release_root'] must be set!"
+
+    release_path = f'{release_root}/dummy'
+
+    if not os.path.isdir(release_path): os.makedirs(release_path)
+
+    with open(f'{release_path}/{branch}-{revision}.txt', 'w') as f: f.write('dummy release file\n')
+
+    return {_StateKey_ : _S_passed_,  _ResultsKey_ : {},  _LogKey_ : f'Config release root set to: {release_root}'}
+
+
+
+def run_python_test(repository_root, working_dir, platform, config):
+
+    import zlib, ssl
+
+    python_environment = local_python_install(platform, config)
+
+    if platform['python'][0] == '2': pass
+    else:
+
+        if platform['os'] == 'mac' and int( platform['python'].split('.')[1] ) > 6 :
+            # SSL certificate test
+            import urllib.request; urllib.request.urlopen('https://benchmark.graylab.jhu.edu')
+
+        ves = [
+            setup_persistent_python_virtual_environment(python_environment, packages='colr dice xdice pdp11games'),
+            setup_python_virtual_environment(working_dir, python_environment, packages='colr dice xdice pdp11games'),
+        ]
+
+        for ve in ves:
+            commands = [
+                'import colr, dice, xdice, pdp11games',
+            ]
+
+            if platform['os'] == 'mac' and int( platform['python'].split('.')[1] ) > 6 :
+                # SSL certificate test
+                commands.append('import urllib.request; urllib.request.urlopen("https://benchmark.graylab.jhu.edu/queue")')
+
+            for command in commands:
+                execute('Testing local Python virtual enviroment...', f"{ve.activate} && {ve.python} -c '{command}'")
+                execute('Testing local Python virtual enviroment...', f"{ve.activate} && python -c '{command}'")
+
+
+
+    return {_StateKey_ : _S_passed_,  _ResultsKey_ : {},  _LogKey_ : f'Done!'}
+
+
+
+def compare(test, results, files_path, previous_results, previous_files_path):
+    """
+    Compare the results of two tests run (new vs. previous) for regression test
+    Take two dict and two paths
+    Must return standard dict with results
+
+    :param test: str
+    :param results: dict
+    :param files_path: str
+    :param previous_results: dict
+    :param previous_files_path: str
+    :rtype: dict
+    """
+    ignore_files = []
+
+    results = dict(tests={}, summary=dict(total=0, failed=0, failed_tests=[]))  # , config={}
+
+    if previous_files_path:
+        for test in os.listdir(files_path):
+            if os.path.isdir(files_path + '/' + test):
+                exclude = ''.join([' --exclude="{}"'.format(f) for f in ignore_files] ) + ' --exclude="*.ignore"'
+                res, brief_diff = execute('Comparing {}...'.format(test), 'diff -rq {exclude} {0}/{test} {1}/{test}'.format(previous_files_path, files_path, test=test, exclude=exclude), return_='tuple')
+                res, full_diff  = execute('Comparing {}...'.format(test), 'diff -r  {exclude} {0}/{test} {1}/{test}'.format(previous_files_path, files_path, test=test, exclude=exclude), return_='tuple')
+                diff = 'Brief Diff:\n' + brief_diff + ( ('\n\nFull Diff:\n' + full_diff[:1024*1024*1]) if full_diff != brief_diff else '' )
+
+                state = _S_failed_ if res else _S_passed_
+                results['tests'][test] = {_StateKey_: state, _LogKey_: diff if state != _S_passed_ else ''}
+
+                results['summary']['total'] += 1
+                if res: results['summary']['failed'] += 1; results['summary']['failed_tests'].append(test)
+
+    else: # no previous tests case, returning 'passed' for all sub_tests
+        for test in os.listdir(files_path):
+            if os.path.isdir(files_path + '/' + test):
+                results['tests'][test] = {_StateKey_: _S_passed_, _LogKey_: 'First run, no previous results available. Skipping comparison...\n'}
+                results['summary']['total'] += 1
+
+    for test in os.listdir(files_path):
+        if os.path.isdir(files_path + '/' + test):
+            if os.path.isfile(files_path+'/'+test+'/.test_did_not_run.log')  or  os.path.isfile(files_path+'/'+test+'/.test_got_timeout_kill.log'):
+                results['tests'][test][_StateKey_] = _S_script_failed_
+                results['tests'][test][_LogKey_] += '\nCompare(...): Marking as "Script failed" due to presense of .test_did_not_run.log or .test_got_timeout_kill.log file!\n'
+                if test not in results['summary']['failed_tests']:
+                    results['summary']['failed'] += 1
+                    results['summary']['failed_tests'].append(test)
+
+    state = _S_failed_ if results['summary']['failed'] else _S_passed_
+
+    return {_StateKey_: state, _LogKey_: 'Comparison dummy log...', _ResultsKey_: results}
+
+
+def run(test, repository_root, working_dir, platform, config, hpc_driver=None, verbose=False, debug=False):
+    if   test == 'state':      return run_state_test      (repository_root, working_dir, platform, config)
+    elif test == 'regression': return run_regression_test (repository_root, working_dir, platform, config)
+    elif test == 'subtests':   return run_subtests_test   (repository_root, working_dir, platform, config)
+    elif test == 'release':    return run_release_test    (repository_root, working_dir, platform, config)
+    elif test == 'python':     return run_python_test     (repository_root, working_dir, platform, config)
+
+    else: raise BenchmarkError(f'Dummy test script does not support run with test={test!r}!')

config/inference/base.yaml

@@ -0,0 +1,136 @@
+# Base inference Configuration.
+
+inference:
+  input_pdb: null
+  num_designs: 10
+  design_startnum: 0
+  ckpt_override_path: null
+  symmetry: null
+  recenter: True
+  radius: 10.0
+  model_only_neighbors: False
+  output_prefix: samples/design
+  write_trajectory: True
+  scaffold_guided: False
+  model_runner: SelfConditioning
+  cautious: True
+  align_motif: True
+  symmetric_self_cond: True
+  final_step: 1
+  deterministic: False
+  trb_save_ckpt_path: null
+  schedule_directory_path: null
+  model_directory_path: null
+
+contigmap:
+  contigs: null
+  inpaint_seq: null
+  provide_seq: null
+  length: null
+
+model:
+  n_extra_block: 4
+  n_main_block: 32
+  n_ref_block: 4
+  d_msa: 256
+  d_msa_full: 64
+  d_pair: 128
+  d_templ: 64
+  n_head_msa: 8
+  n_head_pair: 4
+  n_head_templ: 4
+  d_hidden: 32
+  d_hidden_templ: 32
+  p_drop: 0.15
+  SE3_param_full:
+    num_layers: 1
+    num_channels: 32
+    num_degrees: 2
+    n_heads: 4
+    div: 4
+    l0_in_features: 8
+    l0_out_features: 8
+    l1_in_features: 3
+    l1_out_features: 2
+    num_edge_features: 32
+  SE3_param_topk:
+    num_layers: 1
+    num_channels: 32
+    num_degrees: 2
+    n_heads: 4
+    div: 4
+    l0_in_features: 64
+    l0_out_features: 64
+    l1_in_features: 3
+    l1_out_features: 2
+    num_edge_features: 64
+  freeze_track_motif: False
+  use_motif_timestep: False
+
+diffuser:
+  T: 50
+  b_0: 1e-2
+  b_T: 7e-2
+  schedule_type: linear
+  so3_type: igso3
+  crd_scale: 0.25
+  partial_T: null    
+  so3_schedule_type: linear
+  min_b: 1.5
+  max_b: 2.5
+  min_sigma: 0.02
+  max_sigma: 1.5
+
+denoiser:
+  noise_scale_ca: 1
+  final_noise_scale_ca: 1
+  ca_noise_schedule_type: constant
+  noise_scale_frame: 1
+  final_noise_scale_frame: 1
+  frame_noise_schedule_type: constant
+
+ppi:
+  hotspot_res: null
+
+potentials:
+  guiding_potentials: null 
+  guide_scale: 10
+  guide_decay: constant
+  olig_inter_all : null
+  olig_intra_all : null
+  olig_custom_contact : null
+  substrate: null
+
+contig_settings:
+  ref_idx: null
+  hal_idx: null
+  idx_rf: null
+  inpaint_seq_tensor: null
+
+preprocess:
+  sidechain_input: False
+  motif_sidechain_input: True
+  d_t1d: 22
+  d_t2d: 44
+  prob_self_cond: 0.0
+  str_self_cond: False
+  predict_previous: False
+  
+logging:
+  inputs: False
+
+scaffoldguided:
+  scaffoldguided: False
+  target_pdb: False
+  target_path: null
+  scaffold_list: null
+  scaffold_dir: null
+  sampled_insertion: 0
+  sampled_N: 0
+  sampled_C: 0
+  ss_mask: 0
+  systematic: False
+  target_ss: null
+  target_adj: null
+  mask_loops: True
+  contig_crop: null

config/inference/symmetry.yaml

@@ -0,0 +1,26 @@
+# Config for sampling symmetric assemblies.
+
+defaults:
+  - base
+
+inference:
+  # Symmetry to sample
+  # Available symmetries:
+  # - Cyclic symmetry (C_n) # call as c5
+  # - Dihedral symmetry (D_n) # call as d5
+  # - Tetrahedral symmetry # call as tetrahedral
+  # - Octahedral symmetry # call as octahedral
+  # - Icosahedral symmetry # call as icosahedral
+  symmetry: c2
+
+  # Set to true for computational efficiency 
+  # to avoid memory overhead of modeling all subunits.
+  model_only_neighbors: False
+
+  # Output directory of samples.
+  output_prefix: samples/c2
+
+contigmap:
+  # Specify a single integer value to sample unconditionally.
+  # Must be evenly divisible by the number of chains in the symmetry.
+  contigs: ['100']

docker/Dockerfile

@@ -0,0 +1,50 @@
+# Usage: 
+# git clone https://github.com/RosettaCommons/RFdiffusion.git
+# cd RFdiffusion
+# docker build -f docker/Dockerfile -t rfdiffusion .
+# mkdir $HOME/inputs $HOME/outputs $HOME/models
+# bash scripts/download_models.sh $HOME/models
+# wget -P $HOME/inputs https://files.rcsb.org/view/5TPN.pdb
+
+# docker run -it --rm --gpus all \
+#   -v $HOME/models:$HOME/models \
+#   -v $HOME/inputs:$HOME/inputs \
+#   -v $HOME/outputs:$HOME/outputs \
+#   rfdiffusion \
+#   inference.output_prefix=$HOME/outputs/motifscaffolding \
+#   inference.model_directory_path=$HOME/models \
+#   inference.input_pdb=$HOME/inputs/5TPN.pdb \
+#   inference.num_designs=3 \
+#   'contigmap.contigs=[10-40/A163-181/10-40]'
+
+FROM nvcr.io/nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu20.04
+
+COPY . /app/RFdiffusion/
+
+RUN apt-get -q update \ 
+  && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
+  git \
+  python3.9 \
+  python3-pip \
+  && python3.9 -m pip install -q -U --no-cache-dir pip \
+  && rm -rf /var/lib/apt/lists/* \
+  && apt-get autoremove -y \
+  && apt-get clean \
+  && pip install -q --no-cache-dir \
+  dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html \
+  torch==1.12.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 \
+  e3nn==0.3.3 \
+  wandb==0.12.0 \
+  pynvml==11.0.0 \
+  git+https://github.com/NVIDIA/dllogger#egg=dllogger \
+  decorator==5.1.0 \
+  hydra-core==1.3.2 \
+  pyrsistent==0.19.3 \
+  /app/RFdiffusion/env/SE3Transformer \
+  && pip install --no-cache-dir /app/RFdiffusion --no-deps
+  
+WORKDIR /app/RFdiffusion
+
+ENV DGLBACKEND="pytorch"
+
+ENTRYPOINT ["python3.9", "scripts/run_inference.py"]

env/SE3Transformer/.dockerignore

@@ -0,0 +1,123 @@
+.Trash-0
+.git
+data/
+.DS_Store
+*wandb/
+*.pt
+*.swp
+
+# added by FAFU
+.idea/
+cache/
+downloaded/
+*.lprof
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+**/benchmark
+**/results
+*.pkl
+*.log

env/SE3Transformer/.gitignore

@@ -0,0 +1,121 @@
+data/
+.DS_Store
+*wandb/
+*.pt
+*.swp
+
+# added by FAFU
+.idea/
+cache/
+downloaded/
+*.lprof
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+**/benchmark
+**/results
+*.pkl
+*.log

env/SE3Transformer/Dockerfile

@@ -0,0 +1,58 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+# run docker daemon with --default-runtime=nvidia for GPU detection during build
+# multistage build for DGL with CUDA and FP16
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:21.07-py3
+
+FROM ${FROM_IMAGE_NAME} AS dgl_builder
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update \
+    && apt-get install -y git build-essential python3-dev make cmake \
+    && rm -rf /var/lib/apt/lists/*
+WORKDIR /dgl
+RUN git clone --branch v0.7.0 --recurse-submodules --depth 1 https://github.com/dmlc/dgl.git .
+RUN sed -i 's/"35 50 60 70"/"60 70 80"/g' cmake/modules/CUDA.cmake
+WORKDIR build
+RUN cmake -DUSE_CUDA=ON -DUSE_FP16=ON ..
+RUN make -j8
+
+
+FROM ${FROM_IMAGE_NAME}
+
+RUN rm -rf /workspace/*
+WORKDIR /workspace/se3-transformer
+
+# copy built DGL and install it
+COPY --from=dgl_builder /dgl ./dgl
+RUN cd dgl/python && python setup.py install && cd ../.. && rm -rf dgl
+
+ADD requirements.txt .
+RUN pip install --no-cache-dir --upgrade --pre pip
+RUN pip install --no-cache-dir -r requirements.txt
+ADD . .
+
+ENV DGLBACKEND=pytorch
+ENV OMP_NUM_THREADS=1

env/SE3Transformer/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2021 NVIDIA CORPORATION & AFFILIATES
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

env/SE3Transformer/NOTICE

@@ -0,0 +1,7 @@
+SE(3)-Transformer PyTorch
+
+This repository includes software from https://github.com/FabianFuchsML/se3-transformer-public
+licensed under the MIT License.
+
+This repository includes software from https://github.com/lucidrains/se3-transformer-pytorch
+licensed under the MIT License.

env/SE3Transformer/README.md

@@ -0,0 +1,580 @@
+# SE(3)-Transformers For PyTorch
+
+This repository provides a script and recipe to train the SE(3)-Transformer model to achieve state-of-the-art accuracy. The content of this repository is tested and maintained by NVIDIA.
+
+## Table Of Contents
+- [Model overview](#model-overview)
+    * [Model architecture](#model-architecture)
+    * [Default configuration](#default-configuration)
+    * [Feature support matrix](#feature-support-matrix)
+        * [Features](#features)
+    * [Mixed precision training](#mixed-precision-training)
+        * [Enabling mixed precision](#enabling-mixed-precision)
+          * [Enabling TF32](#enabling-tf32)
+    * [Glossary](#glossary)
+- [Setup](#setup)
+    * [Requirements](#requirements)
+- [Quick Start Guide](#quick-start-guide)
+- [Advanced](#advanced)
+    * [Scripts and sample code](#scripts-and-sample-code)
+    * [Parameters](#parameters)
+    * [Command-line options](#command-line-options)
+    * [Getting the data](#getting-the-data)
+        * [Dataset guidelines](#dataset-guidelines)
+        * [Multi-dataset](#multi-dataset)
+    * [Training process](#training-process)
+    * [Inference process](#inference-process)
+- [Performance](#performance)
+    * [Benchmarking](#benchmarking)
+        * [Training performance benchmark](#training-performance-benchmark)
+        * [Inference performance benchmark](#inference-performance-benchmark)
+    * [Results](#results)
+        * [Training accuracy results](#training-accuracy-results)                         
+            * [Training accuracy: NVIDIA DGX A100 (8x A100 80GB)](#training-accuracy-nvidia-dgx-a100-8x-a100-80gb)  
+            * [Training accuracy: NVIDIA DGX-1 (8x V100 16GB)](#training-accuracy-nvidia-dgx-1-8x-v100-16gb)
+            * [Training stability test](#training-stability-test)
+        * [Training performance results](#training-performance-results)
+            * [Training performance: NVIDIA DGX A100 (8x A100 80GB)](#training-performance-nvidia-dgx-a100-8x-a100-80gb) 
+            * [Training performance: NVIDIA DGX-1 (8x V100 16GB)](#training-performance-nvidia-dgx-1-8x-v100-16gb)
+        * [Inference performance results](#inference-performance-results)
+            * [Inference performance: NVIDIA DGX A100 (1x A100 80GB)](#inference-performance-nvidia-dgx-a100-1x-a100-80gb)
+            * [Inference performance: NVIDIA DGX-1 (1x V100 16GB)](#inference-performance-nvidia-dgx-1-1x-v100-16gb)
+- [Release notes](#release-notes)
+    * [Changelog](#changelog)
+    * [Known issues](#known-issues)
+
+
+
+## Model overview
+
+
+The **SE(3)-Transformer** is a Graph Neural Network using a variant of [self-attention](https://arxiv.org/abs/1706.03762v5) for 3D points and graphs processing.
+This model is [equivariant](https://en.wikipedia.org/wiki/Equivariant_map) under [continuous 3D roto-translations](https://en.wikipedia.org/wiki/Euclidean_group), meaning that when the inputs (graphs or sets of points) rotate in 3D space (or more generally experience a [proper rigid transformation](https://en.wikipedia.org/wiki/Rigid_transformation)), the model outputs either stay invariant or transform with the input.
+A mathematical guarantee of equivariance is important to ensure stable and predictable performance in the presence of nuisance transformations of the data input and when the problem has some inherent symmetries we want to exploit.
+
+
+The model is based on the following publications:
+- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://arxiv.org/abs/2006.10503) (NeurIPS 2020) by Fabian B. Fuchs, Daniel E. Worrall, et al. 
+- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://arxiv.org/abs/1802.08219) by Nathaniel Thomas, Tess Smidt, et al.
+
+A follow-up paper explains how this model can be used iteratively, for example, to predict or refine protein structures:
+
+- [Iterative SE(3)-Transformers](https://arxiv.org/abs/2102.13419) by Fabian B. Fuchs, Daniel E. Worrall, et al. 
+
+Just like [the official implementation](https://github.com/FabianFuchsML/se3-transformer-public), this implementation uses [PyTorch](https://pytorch.org/) and the [Deep Graph Library (DGL)](https://www.dgl.ai/).
+
+The main differences between this implementation of SE(3)-Transformers and the official one are the following:
+
+- Training and inference support for multiple GPUs
+- Training and inference support for [Mixed Precision](https://arxiv.org/abs/1710.03740)
+- The [QM9 dataset from DGL](https://docs.dgl.ai/en/latest/api/python/dgl.data.html#qm9edge-dataset) is used and automatically downloaded
+- Significantly increased throughput
+- Significantly reduced memory consumption
+- The use of layer normalization in the fully connected radial profile layers is an option (`--use_layer_norm`), off by default
+- The use of equivariant normalization between attention layers is an option (`--norm`), off by default
+- The [spherical harmonics](https://en.wikipedia.org/wiki/Spherical_harmonic) and [Clebsch–Gordan coefficients](https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients), used to compute bases matrices, are computed with the [e3nn library](https://e3nn.org/)
+
+
+
+This model enables you to predict quantum chemical properties of small organic molecules in the [QM9 dataset](https://www.nature.com/articles/sdata201422).
+In this case, the exploited symmetry is that these properties do not depend on the orientation or position of the molecules in space.
+
+
+This model is trained with mixed precision using Tensor Cores on NVIDIA Volta, NVIDIA Turing, and the NVIDIA Ampere GPU architectures. Therefore, researchers can get results up to 1.5x faster than training without Tensor Cores while experiencing the benefits of mixed precision training. This model is tested against each NGC monthly container release to ensure consistent accuracy and performance over time.
+
+### Model architecture
+
+The model consists of stacked layers of equivariant graph self-attention and equivariant normalization.
+Lastly, a Tensor Field Network convolution is applied to obtain invariant features. Graph pooling (mean or max over the nodes) is applied to these features, and the result is fed to a final MLP to get scalar predictions.
+
+In this setup, the model is a graph-to-scalar network. The pooling can be removed to obtain a graph-to-graph network, and the final TFN can be modified to output features of any type (invariant scalars, 3D vectors, ...).
+
+
+![Model high-level architecture](./images/se3-transformer.png)
+
+
+### Default configuration
+
+
+SE(3)-Transformers introduce a self-attention layer for graphs that is equivariant to 3D roto-translations. It achieves this by leveraging Tensor Field Networks to build attention weights that are invariant and attention values that are equivariant.
+Combining the equivariant values with the invariant weights gives rise to an equivariant output. This output is normalized while preserving equivariance thanks to equivariant normalization layers operating on feature norms.
+
+
+The following features were implemented in this model:
+
+- Support for edge features of any degree (1D, 3D, 5D, ...), whereas the official implementation only supports scalar invariant edge features (degree 0). Edge features with a degree greater than one are
+concatenated to node features of the same degree. This is required in order to reproduce published results on point cloud processing.
+- Data-parallel multi-GPU training (DDP)
+- Mixed precision training (autocast, gradient scaling)
+- Gradient accumulation
+- Model checkpointing
+
+
+The following performance optimizations were implemented in this model:
+
+
+**General optimizations**
+
+- The option is provided to precompute bases at the beginning of the training instead of computing them at the beginning of each forward pass (`--precompute_bases`)
+- The bases computation is just-in-time (JIT) compiled with `torch.jit.script`
+- The Clebsch-Gordon coefficients are cached in RAM
+
+
+**Tensor Field Network optimizations**
+
+- The last layer of each radial profile network does not add any bias in order to avoid large broadcasting operations
+- The layout (order of dimensions) of the bases tensors is optimized to avoid copies to contiguous memory in the downstream TFN layers
+- When Tensor Cores are available, and the output feature dimension of computed bases is odd, then it is padded with zeros to make more effective use of Tensor Cores (AMP and TF32 precisions)
+- Multiple levels of fusion for TFN convolutions (and radial profiles) are provided and automatically used when conditions are met
+- A low-memory mode is provided that will trade throughput for less memory use (`--low_memory`)
+
+**Self-attention optimizations**
+
+- Attention keys and values are computed by a single partial TFN graph convolution in each attention layer instead of two
+- Graph operations for different output degrees may be fused together if conditions are met
+
+
+**Normalization optimizations**
+
+- The equivariant normalization layer is optimized from multiple layer normalizations to a group normalization on fused norms when certain conditions are met
+    
+
+
+Competitive training results and analysis are provided for the following hyperparameters (identical to the ones in the original publication):
+- Number of layers: 7
+- Number of degrees: 4
+- Number of channels: 32
+- Number of attention heads: 8
+- Channels division: 2
+- Use of equivariant normalization: true
+- Use of layer normalization: true
+- Pooling: max
+
+
+### Feature support matrix
+
+This model supports the following features:: 
+
+| Feature               | SE(3)-Transformer                
+|-----------------------|--------------------------
+|Automatic mixed precision (AMP)   |         Yes 
+|Distributed data parallel (DDP)   |         Yes 
+         
+#### Features
+
+
+**Distributed data parallel (DDP)**
+
+[DistributedDataParallel (DDP)](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) implements data parallelism at the module level that can run across multiple GPUs or machines.
+
+**Automatic Mixed Precision (AMP)**
+
+This implementation uses the native PyTorch AMP implementation of mixed precision training. It allows us to use FP16 training with FP32 master weights by modifying just a few lines of code. A detailed explanation of mixed precision can be found in the next section.
+
+### Mixed precision training
+
+Mixed precision is the combined use of different numerical precisions in a computational method. [Mixed precision](https://arxiv.org/abs/1710.03740) training offers significant computational speedup by performing operations in half-precision format while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of [Tensor Cores](https://developer.nvidia.com/tensor-cores) in NVIDIA Volta, and following with both the NVIDIA Turing and NVIDIA Ampere Architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures. Using [mixed precision training](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html) previously required two steps:
+1.  Porting the model to use the FP16 data type where appropriate.    
+2.  Adding loss scaling to preserve small gradient values.
+
+AMP enables mixed precision training on NVIDIA Volta, NVIDIA Turing, and NVIDIA Ampere GPU architectures automatically. The PyTorch framework code makes all necessary model changes internally.
+
+For information about:
+-   How to train using mixed precision, refer to the [Mixed Precision Training](https://arxiv.org/abs/1710.03740) paper and [Training With Mixed Precision](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html) documentation.
+-   Techniques used for mixed precision training, refer to the [Mixed-Precision Training of Deep Neural Networks](https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/) blog.
+-   APEX tools for mixed precision training, refer to the [NVIDIA Apex: Tools for Easy Mixed-Precision Training in PyTorch](https://devblogs.nvidia.com/apex-pytorch-easy-mixed-precision-training/).
+
+#### Enabling mixed precision
+
+Mixed precision is enabled in PyTorch by using the native [Automatic Mixed Precision package](https://pytorch.org/docs/stable/amp.html), which casts variables to half-precision upon retrieval while storing variables in single-precision format. Furthermore, to preserve small gradient magnitudes in backpropagation, a [loss scaling](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html#lossscaling) step must be included when applying gradients. In PyTorch, loss scaling can be applied automatically using a `GradScaler`.
+Automatic Mixed Precision makes all the adjustments internally in PyTorch, providing two benefits over manual operations. First, programmers need not modify network model code, reducing development and maintenance effort. Second, using AMP maintains forward and backward compatibility with all the APIs for defining and running PyTorch models.
+
+To enable mixed precision, you can simply use the `--amp` flag when running the training or inference scripts.
+
+#### Enabling TF32
+
+TensorFloat-32 (TF32) is the new math mode in [NVIDIA A100](https://www.nvidia.com/en-us/data-center/a100/) GPUs for handling the matrix math, also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on NVIDIA Volta GPUs. 
+
+TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models that require a high dynamic range for weights or activations.
+
+For more information, refer to the [TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x](https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) blog post.
+
+TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.
+
+
+
+### Glossary
+
+**Degree (type)**
+
+In the model, every feature (input, output and hidden) transforms in an equivariant way in relation to the input graph. When we define a feature, we need to choose, in addition to the number of channels, which transformation rule it obeys.
+
+The degree or type of a feature is a positive integer that describes how this feature transforms when the input rotates in 3D.
+
+This is related to [irreducible representations](https://en.wikipedia.org/wiki/Irreducible_representation) of different rotation orders.
+
+The degree of a feature determines its dimensionality. A type-d feature has a dimensionality of 2d+1.
+
+Some common examples include:
+- Degree 0: 1D scalars invariant to rotation
+- Degree 1: 3D vectors that rotate according to 3D rotation matrices
+- Degree 2: 5D vectors that rotate according to 5D [Wigner-D matrices](https://en.wikipedia.org/wiki/Wigner_D-matrix). These can represent symmetric traceless 3x3 matrices.
+
+**Fiber**
+
+A fiber can be viewed as a representation of a set of features of different types or degrees (positive integers), where each feature type transforms according to its rule.
+
+In this repository, a fiber can be seen as a dictionary with degrees as keys and numbers of channels as values.
+
+**Multiplicity**
+
+The multiplicity of a feature of a given type is the number of channels of this feature.
+
+**Tensor Field Network**
+
+A [Tensor Field Network](https://arxiv.org/abs/1802.08219) is a kind of equivariant graph convolution that can combine features of different degrees and produce new ones while preserving equivariance thanks to [tensor products](https://en.wikipedia.org/wiki/Tensor_product).
+
+**Equivariance**
+
+[Equivariance](https://en.wikipedia.org/wiki/Equivariant_map) is a property of a function of model stating that applying a symmetry transformation to the input and then computing the function produces the same result as computing the function and then applying the transformation to the output.
+
+In the case of SE(3)-Transformer, the symmetry group is the group of continuous roto-translations (SE(3)).
+
+## Setup
+
+The following section lists the requirements that you need to meet in order to start training the SE(3)-Transformer model.
+
+### Requirements
+
+This repository contains a Dockerfile which extends the PyTorch 21.07 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+- [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+- PyTorch 21.07+ NGC container
+- Supported GPUs:
+    - [NVIDIA Volta architecture](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/)
+    - [NVIDIA Turing architecture](https://www.nvidia.com/en-us/design-visualization/technologies/turing-architecture/)
+    - [NVIDIA Ampere architecture](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/)
+
+For more information about how to get started with NGC containers, refer to the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:
+- [Getting Started Using NVIDIA GPU Cloud](https://docs.nvidia.com/ngc/ngc-getting-started-guide/index.html)
+- [Accessing And Pulling From The NGC Container Registry](https://docs.nvidia.com/deeplearning/frameworks/user-guide/index.html#accessing_registry)
+- [Running PyTorch](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/running.html#running)
+  
+For those unable to use the PyTorch NGC container to set up the required environment or create your own container, refer to the versioned [NVIDIA Container Support Matrix](https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html).
+
+## Quick Start Guide
+
+To train your model using mixed or TF32 precision with Tensor Cores or FP32, perform the following steps using the default parameters of the SE(3)-Transformer model on the QM9 dataset. For the specifics concerning training and inference, refer to the [Advanced](#advanced) section.
+
+1. Clone the repository.
+    ```
+    git clone https://github.com/NVIDIA/DeepLearningExamples
+    cd DeepLearningExamples/PyTorch/DrugDiscovery/SE3Transformer
+    ```
+   
+2.  Build the `se3-transformer` PyTorch NGC container.
+    ```
+    docker build -t se3-transformer .
+    ```
+
+3.  Start an interactive session in the NGC container to run training/inference.
+    ```
+    mkdir -p results
+    docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/results:/results se3-transformer:latest
+    ```
+
+4. Start training.
+   ```
+   bash scripts/train.sh
+   ```
+
+5. Start inference/predictions.
+   ```
+   bash scripts/predict.sh
+   ```
+
+
+Now that you have your model trained and evaluated, you can choose to compare your training results with our [Training accuracy results](#training-accuracy-results). You can also choose to benchmark your performance to [Training performance benchmark](#training-performance-results) or [Inference performance benchmark](#inference-performance-results). Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the [Results](#results) section.
+
+## Advanced
+
+The following sections provide greater details of the dataset, running training and inference, and the training results.
+
+### Scripts and sample code
+
+In the root directory, the most important files are:
+- `Dockerfile`: container with the basic set of dependencies to run SE(3)-Transformers
+- `requirements.txt`: set of extra requirements to run SE(3)-Transformers
+- `se3_transformer/data_loading/qm9.py`: QM9 data loading and preprocessing, as well as bases precomputation
+- `se3_transformer/model/layers/`: directory containing model architecture layers
+- `se3_transformer/model/transformer.py`: main Transformer module
+- `se3_transformer/model/basis.py`: logic for computing bases matrices
+- `se3_transformer/runtime/training.py`: training script, to be run as a python module
+- `se3_transformer/runtime/inference.py`: inference script, to be run as a python module
+- `se3_transformer/runtime/metrics.py`: MAE metric with support for multi-GPU synchronization
+- `se3_transformer/runtime/loggers.py`: [DLLogger](https://github.com/NVIDIA/dllogger) and [W&B](wandb.ai/) loggers
+
+
+### Parameters
+
+The complete list of the available parameters for the `training.py` script contains:
+
+**General**
+
+- `--epochs`: Number of training epochs (default: `100` for single-GPU)
+- `--batch_size`: Batch size (default: `240`)
+- `--seed`: Set a seed globally (default: `None`)
+- `--num_workers`: Number of dataloading workers (default: `8`)
+- `--amp`: Use Automatic Mixed Precision (default `false`)
+- `--gradient_clip`: Clipping of the gradient norms (default: `None`)
+- `--accumulate_grad_batches`: Gradient accumulation (default: `1`)
+- `--ckpt_interval`: Save a checkpoint every N epochs (default: `-1`)
+- `--eval_interval`: Do an evaluation round every N epochs (default: `1`)
+- `--silent`: Minimize stdout output (default: `false`)
+
+**Paths**
+
+- `--data_dir`: Directory where the data is located or should be downloaded (default: `./data`)
+- `--log_dir`: Directory where the results logs should be saved (default: `/results`)
+- `--save_ckpt_path`: File where the checkpoint should be saved (default: `None`)
+- `--load_ckpt_path`: File of the checkpoint to be loaded (default: `None`)
+
+**Optimizer**
+
+- `--optimizer`: Optimizer to use (default: `adam`)
+- `--learning_rate`: Learning rate to use (default: `0.002` for single-GPU)
+- `--momentum`: Momentum to use (default: `0.9`)
+- `--weight_decay`: Weight decay to use (default: `0.1`)
+
+**QM9 dataset**
+
+- `--task`: Regression task to train on (default: `homo`)
+- `--precompute_bases`: Precompute bases at the beginning of the script during dataset initialization, instead of computing them at the beginning of each forward pass (default: `false`)
+
+**Model architecture**
+
+- `--num_layers`: Number of stacked Transformer layers (default: `7`)
+- `--num_heads`: Number of heads in self-attention (default: `8`)
+- `--channels_div`: Channels division before feeding to attention layer (default: `2`)
+- `--pooling`: Type of graph pooling (default: `max`)
+- `--norm`: Apply a normalization layer after each attention block (default: `false`)
+- `--use_layer_norm`: Apply layer normalization between MLP layers (default: `false`)
+- `--low_memory`: If true, will use fused ops that are slower but use less memory (expect 25 percent less memory). Only has an effect if AMP is enabled on NVIDIA Volta GPUs or if running on Ampere GPUs (default: `false`)
+- `--num_degrees`: Number of degrees to use. Hidden features will have types [0, ..., num_degrees - 1] (default: `4`)
+- `--num_channels`: Number of channels for the hidden features (default: `32`)
+
+
+### Command-line options
+
+To show the full list of available options and their descriptions, use the `-h` or `--help` command-line option, for example: `python -m se3_transformer.runtime.training --help`.
+
+
+### Dataset guidelines
+
+#### Demo dataset
+
+The SE(3)-Transformer was trained on the QM9 dataset.
+
+The QM9 dataset is hosted on DGL servers and downloaded (38MB) automatically when needed. By default, it is stored in the `./data` directory, but this location can be changed with the `--data_dir` argument.
+
+The dataset is saved as a `qm9_edge.npz` file and converted to DGL graphs at runtime.
+
+As input features, we use:
+- Node features (6D):
+    - One-hot-encoded atom type (5D) (atom types: H, C, N, O, F)
+    - Number of protons of each atom (1D)
+- Edge features: one-hot-encoded bond type (4D) (bond types: single, double, triple, aromatic)
+- The relative positions between adjacent nodes (atoms)
+
+#### Custom datasets
+
+To use this network on a new dataset, you can extend the `DataModule` class present in `se3_transformer/data_loading/data_module.py`.
+
+Your custom collate function should return a tuple with:
+
+- A (batched) DGLGraph object
+- A dictionary of node features ({‘{degree}’: tensor})
+- A dictionary of edge features ({‘{degree}’: tensor})
+- (Optional) Precomputed bases as a dictionary
+- Labels as a tensor
+
+You can then modify the `training.py` and `inference.py` scripts to use your new data module.
+
+### Training process
+
+The training script is `se3_transformer/runtime/training.py`, to be run as a module: `python -m se3_transformer.runtime.training`.
+
+**Logs**
+
+By default, the resulting logs are stored in `/results/`. This can be changed with `--log_dir`.
+
+You can connect your existing Weights & Biases account by setting the `WANDB_API_KEY` environment variable.
+
+**Checkpoints**
+
+The argument `--save_ckpt_path` can be set to the path of the file where the checkpoints should be saved.
+`--ckpt_interval` can also be set to the interval (in the number of epochs) between checkpoints.
+
+**Evaluation**
+
+The evaluation metric is the Mean Absolute Error (MAE).
+
+`--eval_interval` can be set to the interval (in the number of epochs) between evaluation rounds. By default, an evaluation round is performed after each epoch.
+
+**Automatic Mixed Precision**
+
+To enable Mixed Precision training, add the `--amp` flag.
+
+**Multi-GPU and multi-node**
+
+The training script supports the PyTorch elastic launcher to run on multiple GPUs or nodes.  Refer to the [official documentation](https://pytorch.org/docs/1.9.0/elastic/run.html).
+
+For example, to train on all available GPUs with AMP:
+
+```
+python -m torch.distributed.run --nnodes=1 --nproc_per_node=gpu --module se3_transformer.runtime.training --amp
+```
+
+
+### Inference process
+
+Inference can be run by using the `se3_transformer.runtime.inference` python module.
+
+The inference script is `se3_transformer/runtime/inference.py`, to be run as a module: `python -m se3_transformer.runtime.inference`.  It requires a pre-trained model checkpoint (to be passed as `--load_ckpt_path`).
+
+
+## Performance
+
+The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to [NVIDIA Data Center Deep Learning Product Performance](https://developer.nvidia.com/deep-learning-performance-training-inference).
+
+### Benchmarking
+
+The following section shows how to run benchmarks measuring the model performance in training and inference modes.
+
+#### Training performance benchmark
+
+To benchmark the training performance on a specific batch size, run `bash scripts/benchmarck_train.sh {BATCH_SIZE}` for single GPU, and `bash scripts/benchmarck_train_multi_gpu.sh {BATCH_SIZE}` for multi-GPU.
+
+#### Inference performance benchmark
+
+To benchmark the inference performance on a specific batch size, run `bash scripts/benchmarck_inference.sh {BATCH_SIZE}`.
+
+### Results
+
+
+The following sections provide details on how we achieved our performance and accuracy in training and inference.
+
+#### Training accuracy results
+
+##### Training accuracy: NVIDIA DGX A100 (8x A100 80GB)
+
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 21.07 NGC container on NVIDIA DGX A100 (8x A100 80GB) GPUs.
+
+| GPUs    | Batch size / GPU    | Absolute error - TF32  | Absolute error - mixed precision  |   Time to train - TF32  |  Time to train - mixed precision | Time to train speedup (mixed precision to TF32) |       
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|  1                 |    240                   |           0.03456                            |        0.03460                                |        1h23min      |    1h03min                |    1.32x              |
+|  8                 |    240                   |           0.03417                            |        0.03424                                |        15min          |    12min                |    1.25x              |
+
+
+##### Training accuracy: NVIDIA DGX-1 (8x V100 16GB)
+
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with (8x V100 16GB) GPUs.
+
+| GPUs    | Batch size / GPU    | Absolute error - FP32  | Absolute error - mixed precision  |   Time to train - FP32  |  Time to train - mixed precision | Time to train speedup (mixed precision to FP32)  |      
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|  1                 |    240                   |           0.03432                            |        0.03439                                |         2h25min         |    1h33min                |    1.56x              |
+|  8                 |    240                   |           0.03380                            |        0.03495                                |        29min          |    20min                |    1.45x              |
+
+
+#### Training performance results
+
+##### Training performance: NVIDIA DGX A100 (8x A100 80GB)
+
+Our results were obtained by running the `scripts/benchmark_train.sh` and `scripts/benchmark_train_multi_gpu.sh` benchmarking scripts in the PyTorch 21.07 NGC container on NVIDIA DGX A100 with 8x A100 80GB GPUs. Performance numbers (in molecules per millisecond) were averaged over five  entire training epochs after a warmup epoch.
+
+| GPUs             | Batch size / GPU     | Throughput - TF32 [mol/ms]                             | Throughput - mixed precision [mol/ms]      | Throughput speedup (mixed precision - TF32)   | Weak scaling - TF32    | Weak scaling - mixed precision |
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|   1              |     240             |   2.21                                       |   2.92                            |   1.32x                         |                      |                                              |
+|   1              |     120              |  1.81                                        |  2.04                             |  1.13x                          |                      |                                              |
+|   8              |     240             |   17.15                                      |     22.95                         |   1.34x                         |   7.76               |    7.86                                     |
+|   8              |     120              |  13.89                                       |    15.62                          |  1.12x                          |       7.67           |    7.66                                       |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+##### Training performance: NVIDIA DGX-1 (8x V100 16GB)
+
+Our results were obtained by running the `scripts/benchmark_train.sh` and `scripts/benchmark_train_multi_gpu.sh` benchmarking scripts in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with 8x V100 16GB GPUs. Performance numbers (in molecules per millisecond) were averaged over five  entire training epochs after a warmup epoch.
+
+| GPUs             | Batch size / GPU     | Throughput - FP32 [mol/ms] | Throughput - mixed precision  [mol/ms]     | Throughput speedup (FP32 - mixed precision)   | Weak scaling - FP32    | Weak scaling - mixed precision |
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|   1              |     240              |    1.25          |    1.88                           |  1.50x                          |                      |                                              |
+|   1              |     120              |    1.03           |   1.41                            |  1.37x                          |                      |                                              |
+|   8              |     240              |    9.33           |   14.02                           |  1.50x                          |      7.46            |      7.46                                    |
+|   8              |     120              |    7.39           |   9.41                           |   1.27x                         |        7.17          |        6.67                                  |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+#### Inference performance results
+
+
+##### Inference performance: NVIDIA DGX A100 (1x A100 80GB)
+
+Our results were obtained by running the `scripts/benchmark_inference.sh` inferencing benchmarking script in the PyTorch 21.07 NGC container on NVIDIA DGX A100 with 1x A100 80GB GPU.
+
+FP16
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 11.60 | 140.94 | 138.29 | 140.12 | 386.40 |
+| 800 | 10.74 | 75.69 | 75.74 | 76.50 | 79.77 |
+| 400 | 8.86 | 45.57 | 46.11 | 46.60 | 49.97 |
+
+TF32
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 8.58 | 189.20 | 186.39 | 187.71 | 420.28 |
+| 800 | 8.28 | 97.56 | 97.20 | 97.73 | 101.13 |
+| 400 | 7.55 | 53.38 | 53.72 | 54.48 | 56.62 |
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+
+##### Inference performance: NVIDIA DGX-1 (1x V100 16GB)
+
+Our results were obtained by running the `scripts/benchmark_inference.sh` inferencing benchmarking script in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with 1x V100 16GB GPU.
+
+FP16
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 6.42 | 254.54 | 247.97 | 249.29 | 721.15 |
+| 800 | 6.13 | 132.07 | 131.90 | 132.70 | 140.15 |
+| 400 | 5.37 | 75.12 | 76.01 | 76.66 | 79.90 |
+
+FP32
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 3.39 | 475.86 | 473.82 | 475.64 | 891.18 |
+| 800 | 3.36 | 239.17 | 240.64 | 241.65 | 243.70 |
+| 400 | 3.17 | 126.67 | 128.19 | 128.82 | 130.54 |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+## Release notes
+
+### Changelog
+
+August 2021
+- Initial release
+
+### Known issues
+
+If you encounter `OSError: [Errno 12] Cannot allocate memory` during the Dataloader iterator creation (more precisely during the `fork()`, this is most likely due to the use of the `--precompute_bases` flag. If you cannot add more RAM or Swap to your machine, it is recommended to turn off bases precomputation by removing the `--precompute_bases` flag or using `--precompute_bases false`.

env/SE3Transformer/build/lib/se3_transformer/data_loading/__init__.py

@@ -0,0 +1 @@
+from .qm9 import QM9DataModule

env/SE3Transformer/build/lib/se3_transformer/data_loading/data_module.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import torch.distributed as dist
+from abc import ABC
+from torch.utils.data import DataLoader, DistributedSampler, Dataset
+
+from se3_transformer.runtime.utils import get_local_rank
+
+
+def _get_dataloader(dataset: Dataset, shuffle: bool, **kwargs) -> DataLoader:
+    # Classic or distributed dataloader depending on the context
+    sampler = DistributedSampler(dataset, shuffle=shuffle) if dist.is_initialized() else None
+    return DataLoader(dataset, shuffle=(shuffle and sampler is None), sampler=sampler, **kwargs)
+
+
+class DataModule(ABC):
+    """ Abstract DataModule. Children must define self.ds_{train | val | test}. """
+
+    def __init__(self, **dataloader_kwargs):
+        super().__init__()
+        if get_local_rank() == 0:
+            self.prepare_data()
+
+        # Wait until rank zero has prepared the data (download, preprocessing, ...)
+        if dist.is_initialized():
+            dist.barrier(device_ids=[get_local_rank()])
+
+        self.dataloader_kwargs = {'pin_memory': True, 'persistent_workers': True, **dataloader_kwargs}
+        self.ds_train, self.ds_val, self.ds_test = None, None, None
+
+    def prepare_data(self):
+        """ Method called only once per node. Put here any downloading or preprocessing """
+        pass
+
+    def train_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_train, shuffle=True, **self.dataloader_kwargs)
+
+    def val_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_val, shuffle=False, **self.dataloader_kwargs)
+
+    def test_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_test, shuffle=False, **self.dataloader_kwargs)

env/SE3Transformer/build/lib/se3_transformer/data_loading/qm9.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+from typing import Tuple
+
+import dgl
+import pathlib
+import torch
+from dgl.data import QM9EdgeDataset
+from dgl import DGLGraph
+from torch import Tensor
+from torch.utils.data import random_split, DataLoader, Dataset
+from tqdm import tqdm
+
+from se3_transformer.data_loading.data_module import DataModule
+from se3_transformer.model.basis import get_basis
+from se3_transformer.runtime.utils import get_local_rank, str2bool, using_tensor_cores
+
+
+def _get_relative_pos(qm9_graph: DGLGraph) -> Tensor:
+    x = qm9_graph.ndata['pos']
+    src, dst = qm9_graph.edges()
+    rel_pos = x[dst] - x[src]
+    return rel_pos
+
+
+def _get_split_sizes(full_dataset: Dataset) -> Tuple[int, int, int]:
+    len_full = len(full_dataset)
+    len_train = 100_000
+    len_test = int(0.1 * len_full)
+    len_val = len_full - len_train - len_test
+    return len_train, len_val, len_test
+
+
+class QM9DataModule(DataModule):
+    """
+    Datamodule wrapping https://docs.dgl.ai/en/latest/api/python/dgl.data.html#qm9edge-dataset
+    Training set is 100k molecules. Test set is 10% of the dataset. Validation set is the rest.
+    This includes all the molecules from QM9 except the ones that are uncharacterized.
+    """
+
+    NODE_FEATURE_DIM = 6
+    EDGE_FEATURE_DIM = 4
+
+    def __init__(self,
+                 data_dir: pathlib.Path,
+                 task: str = 'homo',
+                 batch_size: int = 240,
+                 num_workers: int = 8,
+                 num_degrees: int = 4,
+                 amp: bool = False,
+                 precompute_bases: bool = False,
+                 **kwargs):
+        self.data_dir = data_dir  # This needs to be before __init__ so that prepare_data has access to it
+        super().__init__(batch_size=batch_size, num_workers=num_workers, collate_fn=self._collate)
+        self.amp = amp
+        self.task = task
+        self.batch_size = batch_size
+        self.num_degrees = num_degrees
+
+        qm9_kwargs = dict(label_keys=[self.task], verbose=False, raw_dir=str(data_dir))
+        if precompute_bases:
+            bases_kwargs = dict(max_degree=num_degrees - 1, use_pad_trick=using_tensor_cores(amp), amp=amp)
+            full_dataset = CachedBasesQM9EdgeDataset(bases_kwargs=bases_kwargs, batch_size=batch_size,
+                                                     num_workers=num_workers, **qm9_kwargs)
+        else:
+            full_dataset = QM9EdgeDataset(**qm9_kwargs)
+
+        self.ds_train, self.ds_val, self.ds_test = random_split(full_dataset, _get_split_sizes(full_dataset),
+                                                                generator=torch.Generator().manual_seed(0))
+
+        train_targets = full_dataset.targets[self.ds_train.indices, full_dataset.label_keys[0]]
+        self.targets_mean = train_targets.mean()
+        self.targets_std = train_targets.std()
+
+    def prepare_data(self):
+        # Download the QM9 preprocessed data
+        QM9EdgeDataset(verbose=True, raw_dir=str(self.data_dir))
+
+    def _collate(self, samples):
+        graphs, y, *bases = map(list, zip(*samples))
+        batched_graph = dgl.batch(graphs)
+        edge_feats = {'0': batched_graph.edata['edge_attr'][..., None]}
+        batched_graph.edata['rel_pos'] = _get_relative_pos(batched_graph)
+        # get node features
+        node_feats = {'0': batched_graph.ndata['attr'][:, :6, None]}
+        targets = (torch.cat(y) - self.targets_mean) / self.targets_std
+
+        if bases:
+            # collate bases
+            all_bases = {
+                key: torch.cat([b[key] for b in bases[0]], dim=0)
+                for key in bases[0][0].keys()
+            }
+
+            return batched_graph, node_feats, edge_feats, all_bases, targets
+        else:
+            return batched_graph, node_feats, edge_feats, targets
+
+    @staticmethod
+    def add_argparse_args(parent_parser):
+        parser = parent_parser.add_argument_group("QM9 dataset")
+        parser.add_argument('--task', type=str, default='homo', const='homo', nargs='?',
+                            choices=['mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'U0', 'U', 'H', 'G', 'Cv',
+                                     'U0_atom', 'U_atom', 'H_atom', 'G_atom', 'A', 'B', 'C'],
+                            help='Regression task to train on')
+        parser.add_argument('--precompute_bases', type=str2bool, nargs='?', const=True, default=False,
+                            help='Precompute bases at the beginning of the script during dataset initialization,'
+                                 ' instead of computing them at the beginning of each forward pass.')
+        return parent_parser
+
+    def __repr__(self):
+        return f'QM9({self.task})'
+
+
+class CachedBasesQM9EdgeDataset(QM9EdgeDataset):
+    """ Dataset extending the QM9 dataset from DGL with precomputed (cached in RAM) pairwise bases """
+
+    def __init__(self, bases_kwargs: dict, batch_size: int, num_workers: int, *args, **kwargs):
+        """
+        :param bases_kwargs:  Arguments to feed the bases computation function
+        :param batch_size:    Batch size to use when iterating over the dataset for computing bases
+        """
+        self.bases_kwargs = bases_kwargs
+        self.batch_size = batch_size
+        self.bases = None
+        self.num_workers = num_workers
+        super().__init__(*args, **kwargs)
+
+    def load(self):
+        super().load()
+        # Iterate through the dataset and compute bases (pairwise only)
+        # Potential improvement: use multi-GPU and gather
+        dataloader = DataLoader(self, shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers,
+                                collate_fn=lambda samples: dgl.batch([sample[0] for sample in samples]))
+        bases = []
+        for i, graph in tqdm(enumerate(dataloader), total=len(dataloader), desc='Precomputing QM9 bases',
+                             disable=get_local_rank() != 0):
+            rel_pos = _get_relative_pos(graph)
+            # Compute the bases with the GPU but convert the result to CPU to store in RAM
+            bases.append({k: v.cpu() for k, v in get_basis(rel_pos.cuda(), **self.bases_kwargs).items()})
+        self.bases = bases  # Assign at the end so that __getitem__ isn't confused
+
+    def __getitem__(self, idx: int):
+        graph, label = super().__getitem__(idx)
+
+        if self.bases:
+            bases_idx = idx // self.batch_size
+            bases_cumsum_idx = self.ne_cumsum[idx] - self.ne_cumsum[bases_idx * self.batch_size]
+            bases_cumsum_next_idx = self.ne_cumsum[idx + 1] - self.ne_cumsum[bases_idx * self.batch_size]
+            return graph, label, {key: basis[bases_cumsum_idx:bases_cumsum_next_idx] for key, basis in
+                                  self.bases[bases_idx].items()}
+        else:
+            return graph, label

env/SE3Transformer/build/lib/se3_transformer/model/__init__.py

@@ -0,0 +1,2 @@
+from .transformer import SE3Transformer, SE3TransformerPooled
+from .fiber import Fiber

env/SE3Transformer/build/lib/se3_transformer/model/basis.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from functools import lru_cache
+from typing import Dict, List
+
+import e3nn.o3 as o3
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.runtime.utils import degree_to_dim
+
+
+@lru_cache(maxsize=None)
+def get_clebsch_gordon(J: int, d_in: int, d_out: int, device) -> Tensor:
+    """ Get the (cached) Q^{d_out,d_in}_J matrices from equation (8) """
+    return o3.wigner_3j(J, d_in, d_out, dtype=torch.float64, device=device).permute(2, 1, 0)
+
+
+@lru_cache(maxsize=None)
+def get_all_clebsch_gordon(max_degree: int, device) -> List[List[Tensor]]:
+    all_cb = []
+    for d_in in range(max_degree + 1):
+        for d_out in range(max_degree + 1):
+            K_Js = []
+            for J in range(abs(d_in - d_out), d_in + d_out + 1):
+                K_Js.append(get_clebsch_gordon(J, d_in, d_out, device))
+            all_cb.append(K_Js)
+    return all_cb
+
+
+def get_spherical_harmonics(relative_pos: Tensor, max_degree: int) -> List[Tensor]:
+    all_degrees = list(range(2 * max_degree + 1))
+    with nvtx_range('spherical harmonics'):
+        sh = o3.spherical_harmonics(all_degrees, relative_pos, normalize=True)
+        return torch.split(sh, [degree_to_dim(d) for d in all_degrees], dim=1)
+
+
+@torch.jit.script
+def get_basis_script(max_degree: int,
+                     use_pad_trick: bool,
+                     spherical_harmonics: List[Tensor],
+                     clebsch_gordon: List[List[Tensor]],
+                     amp: bool) -> Dict[str, Tensor]:
+    """
+    Compute pairwise bases matrices for degrees up to max_degree
+    :param max_degree:            Maximum input or output degree
+    :param use_pad_trick:         Pad some of the odd dimensions for a better use of Tensor Cores
+    :param spherical_harmonics:   List of computed spherical harmonics
+    :param clebsch_gordon:        List of computed CB-coefficients
+    :param amp:                   When true, return bases in FP16 precision
+    """
+    basis = {}
+    idx = 0
+    # Double for loop instead of product() because of JIT script
+    for d_in in range(max_degree + 1):
+        for d_out in range(max_degree + 1):
+            key = f'{d_in},{d_out}'
+            K_Js = []
+            for freq_idx, J in enumerate(range(abs(d_in - d_out), d_in + d_out + 1)):
+                Q_J = clebsch_gordon[idx][freq_idx]
+                K_Js.append(torch.einsum('n f, k l f -> n l k', spherical_harmonics[J].float(), Q_J.float()))
+
+            basis[key] = torch.stack(K_Js, 2)  # Stack on second dim so order is n l f k
+            if amp:
+                basis[key] = basis[key].half()
+            if use_pad_trick:
+                basis[key] = F.pad(basis[key], (0, 1))  # Pad the k dimension, that can be sliced later
+
+            idx += 1
+
+    return basis
+
+
+@torch.jit.script
+def update_basis_with_fused(basis: Dict[str, Tensor],
+                            max_degree: int,
+                            use_pad_trick: bool,
+                            fully_fused: bool) -> Dict[str, Tensor]:
+    """ Update the basis dict with partially and optionally fully fused bases """
+    num_edges = basis['0,0'].shape[0]
+    device = basis['0,0'].device
+    dtype = basis['0,0'].dtype
+    sum_dim = sum([degree_to_dim(d) for d in range(max_degree + 1)])
+
+    # Fused per output degree
+    for d_out in range(max_degree + 1):
+        sum_freq = sum([degree_to_dim(min(d, d_out)) for d in range(max_degree + 1)])
+        basis_fused = torch.zeros(num_edges, sum_dim, sum_freq, degree_to_dim(d_out) + int(use_pad_trick),
+                                  device=device, dtype=dtype)
+        acc_d, acc_f = 0, 0
+        for d_in in range(max_degree + 1):
+            basis_fused[:, acc_d:acc_d + degree_to_dim(d_in), acc_f:acc_f + degree_to_dim(min(d_out, d_in)),
+            :degree_to_dim(d_out)] = basis[f'{d_in},{d_out}'][:, :, :, :degree_to_dim(d_out)]
+
+            acc_d += degree_to_dim(d_in)
+            acc_f += degree_to_dim(min(d_out, d_in))
+
+        basis[f'out{d_out}_fused'] = basis_fused
+
+    # Fused per input degree
+    for d_in in range(max_degree + 1):
+        sum_freq = sum([degree_to_dim(min(d, d_in)) for d in range(max_degree + 1)])
+        basis_fused = torch.zeros(num_edges, degree_to_dim(d_in), sum_freq, sum_dim,
+                                  device=device, dtype=dtype)
+        acc_d, acc_f = 0, 0
+        for d_out in range(max_degree + 1):
+            basis_fused[:, :, acc_f:acc_f + degree_to_dim(min(d_out, d_in)), acc_d:acc_d + degree_to_dim(d_out)] \
+                = basis[f'{d_in},{d_out}'][:, :, :, :degree_to_dim(d_out)]
+
+            acc_d += degree_to_dim(d_out)
+            acc_f += degree_to_dim(min(d_out, d_in))
+
+        basis[f'in{d_in}_fused'] = basis_fused
+
+    if fully_fused:
+        # Fully fused
+        # Double sum this way because of JIT script
+        sum_freq = sum([
+            sum([degree_to_dim(min(d_in, d_out)) for d_in in range(max_degree + 1)]) for d_out in range(max_degree + 1)
+        ])
+        basis_fused = torch.zeros(num_edges, sum_dim, sum_freq, sum_dim, device=device, dtype=dtype)
+
+        acc_d, acc_f = 0, 0
+        for d_out in range(max_degree + 1):
+            b = basis[f'out{d_out}_fused']
+            basis_fused[:, :, acc_f:acc_f + b.shape[2], acc_d:acc_d + degree_to_dim(d_out)] = b[:, :, :,
+                                                                                              :degree_to_dim(d_out)]
+            acc_f += b.shape[2]
+            acc_d += degree_to_dim(d_out)
+
+        basis['fully_fused'] = basis_fused
+
+    del basis['0,0']  # We know that the basis for l = k = 0 is filled with a constant
+    return basis
+
+
+def get_basis(relative_pos: Tensor,
+              max_degree: int = 4,
+              compute_gradients: bool = False,
+              use_pad_trick: bool = False,
+              amp: bool = False) -> Dict[str, Tensor]:
+    with nvtx_range('spherical harmonics'):
+        spherical_harmonics = get_spherical_harmonics(relative_pos, max_degree)
+    with nvtx_range('CB coefficients'):
+        clebsch_gordon = get_all_clebsch_gordon(max_degree, relative_pos.device)
+
+    with torch.autograd.set_grad_enabled(compute_gradients):
+        with nvtx_range('bases'):
+            basis = get_basis_script(max_degree=max_degree,
+                                     use_pad_trick=use_pad_trick,
+                                     spherical_harmonics=spherical_harmonics,
+                                     clebsch_gordon=clebsch_gordon,
+                                     amp=amp)
+            return basis

env/SE3Transformer/build/lib/se3_transformer/model/fiber.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from collections import namedtuple
+from itertools import product
+from typing import Dict
+
+import torch
+from torch import Tensor
+
+from se3_transformer.runtime.utils import degree_to_dim
+
+FiberEl = namedtuple('FiberEl', ['degree', 'channels'])
+
+
+class Fiber(dict):
+    """
+    Describes the structure of some set of features.
+    Features are split into types (0, 1, 2, 3, ...). A feature of type k has a dimension of 2k+1.
+    Type-0 features: invariant scalars
+    Type-1 features: equivariant 3D vectors
+    Type-2 features: equivariant symmetric traceless matrices
+    ...
+
+    As inputs to a SE3 layer, there can be many features of the same types, and many features of different types.
+    The 'multiplicity' or 'number of channels' is the number of features of a given type.
+    This class puts together all the degrees and their multiplicities in order to describe
+        the inputs, outputs or hidden features of SE3 layers.
+    """
+
+    def __init__(self, structure):
+        if isinstance(structure, dict):
+            structure = [FiberEl(int(d), int(m)) for d, m in sorted(structure.items(), key=lambda x: x[1])]
+        elif not isinstance(structure[0], FiberEl):
+            structure = list(map(lambda t: FiberEl(*t), sorted(structure, key=lambda x: x[1])))
+        self.structure = structure
+        super().__init__({d: m for d, m in self.structure})
+
+    @property
+    def degrees(self):
+        return sorted([t.degree for t in self.structure])
+
+    @property
+    def channels(self):
+        return [self[d] for d in self.degrees]
+
+    @property
+    def num_features(self):
+        """ Size of the resulting tensor if all features were concatenated together """
+        return sum(t.channels * degree_to_dim(t.degree) for t in self.structure)
+
+    @staticmethod
+    def create(num_degrees: int, num_channels: int):
+        """ Create a Fiber with degrees 0..num_degrees-1, all with the same multiplicity """
+        return Fiber([(degree, num_channels) for degree in range(num_degrees)])
+
+    @staticmethod
+    def from_features(feats: Dict[str, Tensor]):
+        """ Infer the Fiber structure from a feature dict """
+        structure = {}
+        for k, v in feats.items():
+            degree = int(k)
+            assert len(v.shape) == 3, 'Feature shape should be (N, C, 2D+1)'
+            assert v.shape[-1] == degree_to_dim(degree)
+            structure[degree] = v.shape[-2]
+        return Fiber(structure)
+
+    def __getitem__(self, degree: int):
+        """ fiber[degree] returns the multiplicity for this degree """
+        return dict(self.structure).get(degree, 0)
+
+    def __iter__(self):
+        """ Iterate over namedtuples (degree, channels) """
+        return iter(self.structure)
+
+    def __mul__(self, other):
+        """
+        If other in an int, multiplies all the multiplicities by other.
+        If other is a fiber, returns the cartesian product.
+        """
+        if isinstance(other, Fiber):
+            return product(self.structure, other.structure)
+        elif isinstance(other, int):
+            return Fiber({t.degree: t.channels * other for t in self.structure})
+
+    def __add__(self, other):
+        """
+        If other in an int, add other to all the multiplicities.
+        If other is a fiber, add the multiplicities of the fibers together.
+        """
+        if isinstance(other, Fiber):
+            return Fiber({t.degree: t.channels + other[t.degree] for t in self.structure})
+        elif isinstance(other, int):
+            return Fiber({t.degree: t.channels + other for t in self.structure})
+
+    def __repr__(self):
+        return str(self.structure)
+
+    @staticmethod
+    def combine_max(f1, f2):
+        """ Combine two fiber by taking the maximum multiplicity for each degree in both fibers """
+        new_dict = dict(f1.structure)
+        for k, m in f2.structure:
+            new_dict[k] = max(new_dict.get(k, 0), m)
+
+        return Fiber(list(new_dict.items()))
+
+    @staticmethod
+    def combine_selectively(f1, f2):
+        """ Combine two fiber by taking the sum of multiplicities for each degree in the first fiber """
+        # only use orders which occur in fiber f1
+        new_dict = dict(f1.structure)
+        for k in f1.degrees:
+            if k in f2.degrees:
+                new_dict[k] += f2[k]
+        return Fiber(list(new_dict.items()))
+
+    def to_attention_heads(self, tensors: Dict[str, Tensor], num_heads: int):
+        # dict(N, num_channels, 2d+1) -> (N, num_heads, -1)
+        fibers = [tensors[str(degree)].reshape(*tensors[str(degree)].shape[:-2], num_heads, -1) for degree in
+                  self.degrees]
+        fibers = torch.cat(fibers, -1)
+        return fibers

env/SE3Transformer/build/lib/se3_transformer/model/layers/__init__.py

@@ -0,0 +1,5 @@
+from .linear import LinearSE3
+from .norm import NormSE3
+from .pooling import GPooling
+from .convolution import ConvSE3
+from .attention import AttentionBlockSE3

env/SE3Transformer/build/lib/se3_transformer/model/layers/attention.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import dgl
+import numpy as np
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from dgl.ops import edge_softmax
+from torch import Tensor
+from typing import Dict, Optional, Union
+
+from se3_transformer.model.fiber import Fiber
+from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel
+from se3_transformer.model.layers.linear import LinearSE3
+from se3_transformer.runtime.utils import degree_to_dim, aggregate_residual, unfuse_features
+from torch.cuda.nvtx import range as nvtx_range
+
+
+class AttentionSE3(nn.Module):
+    """ Multi-headed sparse graph self-attention (SE(3)-equivariant) """
+
+    def __init__(
+            self,
+            num_heads: int,
+            key_fiber: Fiber,
+            value_fiber: Fiber
+    ):
+        """
+        :param num_heads:     Number of attention heads
+        :param key_fiber:     Fiber for the keys (and also for the queries)
+        :param value_fiber:   Fiber for the values
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.key_fiber = key_fiber
+        self.value_fiber = value_fiber
+
+    def forward(
+            self,
+            value: Union[Tensor, Dict[str, Tensor]],  # edge features (may be fused)
+            key: Union[Tensor, Dict[str, Tensor]],  # edge features (may be fused)
+            query: Dict[str, Tensor],  # node features
+            graph: DGLGraph
+    ):
+        with nvtx_range('AttentionSE3'):
+            with nvtx_range('reshape keys and queries'):
+                if isinstance(key, Tensor):
+                    # case where features of all types are fused
+                    key = key.reshape(key.shape[0], self.num_heads, -1)
+                    # need to reshape queries that way to keep the same layout as keys
+                    out = torch.cat([query[str(d)] for d in self.key_fiber.degrees], dim=-1)
+                    query = out.reshape(list(query.values())[0].shape[0], self.num_heads, -1)
+                else:
+                    # features are not fused, need to fuse and reshape them
+                    key = self.key_fiber.to_attention_heads(key, self.num_heads)
+                    query = self.key_fiber.to_attention_heads(query, self.num_heads)
+
+            with nvtx_range('attention dot product + softmax'):
+                # Compute attention weights (softmax of inner product between key and query)
+                edge_weights = dgl.ops.e_dot_v(graph, key, query).squeeze(-1)
+                edge_weights /= np.sqrt(self.key_fiber.num_features)
+                edge_weights = edge_softmax(graph, edge_weights)
+                edge_weights = edge_weights[..., None, None]
+
+            with nvtx_range('weighted sum'):
+                if isinstance(value, Tensor):
+                    # features of all types are fused
+                    v = value.view(value.shape[0], self.num_heads, -1, value.shape[-1])
+                    weights = edge_weights * v
+                    feat_out = dgl.ops.copy_e_sum(graph, weights)
+                    feat_out = feat_out.view(feat_out.shape[0], -1, feat_out.shape[-1])  # merge heads
+                    out = unfuse_features(feat_out, self.value_fiber.degrees)
+                else:
+                    out = {}
+                    for degree, channels in self.value_fiber:
+                        v = value[str(degree)].view(-1, self.num_heads, channels // self.num_heads,
+                                                    degree_to_dim(degree))
+                        weights = edge_weights * v
+                        res = dgl.ops.copy_e_sum(graph, weights)
+                        out[str(degree)] = res.view(-1, channels, degree_to_dim(degree))  # merge heads
+
+                return out
+
+
+class AttentionBlockSE3(nn.Module):
+    """ Multi-headed sparse graph self-attention block with skip connection, linear projection (SE(3)-equivariant) """
+
+    def __init__(
+            self,
+            fiber_in: Fiber,
+            fiber_out: Fiber,
+            fiber_edge: Optional[Fiber] = None,
+            num_heads: int = 4,
+            channels_div: int = 2,
+            use_layer_norm: bool = False,
+            max_degree: bool = 4,
+            fuse_level: ConvSE3FuseLevel = ConvSE3FuseLevel.FULL,
+            **kwargs
+    ):
+        """
+        :param fiber_in:         Fiber describing the input features
+        :param fiber_out:        Fiber describing the output features
+        :param fiber_edge:       Fiber describing the edge features (node distances excluded)
+        :param num_heads:        Number of attention heads
+        :param channels_div:     Divide the channels by this integer for computing values
+        :param use_layer_norm:   Apply layer normalization between MLP layers
+        :param max_degree:       Maximum degree used in the bases computation
+        :param fuse_level:       Maximum fuse level to use in TFN convolutions
+        """
+        super().__init__()
+        if fiber_edge is None:
+            fiber_edge = Fiber({})
+        self.fiber_in = fiber_in
+        # value_fiber has same structure as fiber_out but #channels divided by 'channels_div'
+        value_fiber = Fiber([(degree, channels // channels_div) for degree, channels in fiber_out])
+        # key_query_fiber has the same structure as fiber_out, but only degrees which are in in_fiber
+        # (queries are merely projected, hence degrees have to match input)
+        key_query_fiber = Fiber([(fe.degree, fe.channels) for fe in value_fiber if fe.degree in fiber_in.degrees])
+
+        self.to_key_value = ConvSE3(fiber_in, value_fiber + key_query_fiber, pool=False, fiber_edge=fiber_edge,
+                                    use_layer_norm=use_layer_norm, max_degree=max_degree, fuse_level=fuse_level,
+                                    allow_fused_output=True)
+        self.to_query = LinearSE3(fiber_in, key_query_fiber)
+        self.attention = AttentionSE3(num_heads, key_query_fiber, value_fiber)
+        self.project = LinearSE3(value_fiber + fiber_in, fiber_out)
+
+    def forward(
+            self,
+            node_features: Dict[str, Tensor],
+            edge_features: Dict[str, Tensor],
+            graph: DGLGraph,
+            basis: Dict[str, Tensor]
+    ):
+        with nvtx_range('AttentionBlockSE3'):
+            with nvtx_range('keys / values'):
+                fused_key_value = self.to_key_value(node_features, edge_features, graph, basis)
+                key, value = self._get_key_value_from_fused(fused_key_value)
+
+            with nvtx_range('queries'):
+                query = self.to_query(node_features)
+            
+            z = self.attention(value, key, query, graph)
+            z_concat = aggregate_residual(node_features, z, 'cat')
+            return self.project(z_concat)
+
+    def _get_key_value_from_fused(self, fused_key_value):
+        # Extract keys and queries features from fused features
+        if isinstance(fused_key_value, Tensor):
+            # Previous layer was a fully fused convolution
+            value, key = torch.chunk(fused_key_value, chunks=2, dim=-2)
+        else:
+            key, value = {}, {}
+            for degree, feat in fused_key_value.items():
+                if int(degree) in self.fiber_in.degrees:
+                    value[degree], key[degree] = torch.chunk(feat, chunks=2, dim=-2)
+                else:
+                    value[degree] = feat
+
+        return key, value

env/SE3Transformer/build/lib/se3_transformer/model/layers/convolution.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from enum import Enum
+from itertools import product
+from typing import Dict
+
+import dgl
+import numpy as np
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.model.fiber import Fiber
+from se3_transformer.runtime.utils import degree_to_dim, unfuse_features
+
+
+class ConvSE3FuseLevel(Enum):
+    """
+    Enum to select a maximum level of fusing optimizations that will be applied when certain conditions are met.
+    If a desired level L is picked and the level L cannot be applied to a level, other fused ops < L are considered.
+    A higher level means faster training, but also more memory usage.
+    If you are tight on memory and want to feed large inputs to the network, choose a low value.
+    If you want to train fast, choose a high value.
+    Recommended value is FULL with AMP.
+
+    Fully fused TFN convolutions requirements:
+    - all input channels are the same
+    - all output channels are the same
+    - input degrees span the range [0, ..., max_degree]
+    - output degrees span the range [0, ..., max_degree]
+
+    Partially fused TFN convolutions requirements:
+    * For fusing by output degree:
+    - all input channels are the same
+    - input degrees span the range [0, ..., max_degree]
+    * For fusing by input degree:
+    - all output channels are the same
+    - output degrees span the range [0, ..., max_degree]
+
+    Original TFN pairwise convolutions: no requirements
+    """
+
+    FULL = 2
+    PARTIAL = 1
+    NONE = 0
+
+
+class RadialProfile(nn.Module):
+    """
+    Radial profile function.
+    Outputs weights used to weigh basis matrices in order to get convolution kernels.
+    In TFN notation: $R^{l,k}$
+    In SE(3)-Transformer notation: $\phi^{l,k}$
+
+    Note:
+        In the original papers, this function only depends on relative node distances ||x||.
+        Here, we allow this function to also take as input additional invariant edge features.
+        This does not break equivariance and adds expressive power to the model.
+
+    Diagram:
+        invariant edge features (node distances included) ───> MLP layer (shared across edges) ───> radial weights
+    """
+
+    def __init__(
+            self,
+            num_freq: int,
+            channels_in: int,
+            channels_out: int,
+            edge_dim: int = 1,
+            mid_dim: int = 32,
+            use_layer_norm: bool = False
+    ):
+        """
+        :param num_freq:         Number of frequencies
+        :param channels_in:      Number of input channels
+        :param channels_out:     Number of output channels
+        :param edge_dim:         Number of invariant edge features (input to the radial function)
+        :param mid_dim:          Size of the hidden MLP layers
+        :param use_layer_norm:   Apply layer normalization between MLP layers
+        """
+        super().__init__()
+        modules = [
+            nn.Linear(edge_dim, mid_dim),
+            nn.LayerNorm(mid_dim) if use_layer_norm else None,
+            nn.ReLU(),
+            nn.Linear(mid_dim, mid_dim),
+            nn.LayerNorm(mid_dim) if use_layer_norm else None,
+            nn.ReLU(),
+            nn.Linear(mid_dim, num_freq * channels_in * channels_out, bias=False)
+        ]
+
+        self.net = nn.Sequential(*[m for m in modules if m is not None])
+
+    def forward(self, features: Tensor) -> Tensor:
+        return self.net(features)
+
+
+class VersatileConvSE3(nn.Module):
+    """
+    Building block for TFN convolutions.
+    This single module can be used for fully fused convolutions, partially fused convolutions, or pairwise convolutions.
+    """
+
+    def __init__(self,
+                 freq_sum: int,
+                 channels_in: int,
+                 channels_out: int,
+                 edge_dim: int,
+                 use_layer_norm: bool,
+                 fuse_level: ConvSE3FuseLevel):
+        super().__init__()
+        self.freq_sum = freq_sum
+        self.channels_out = channels_out
+        self.channels_in = channels_in
+        self.fuse_level = fuse_level
+        self.radial_func = RadialProfile(num_freq=freq_sum,
+                                         channels_in=channels_in,
+                                         channels_out=channels_out,
+                                         edge_dim=edge_dim,
+                                         use_layer_norm=use_layer_norm)
+
+    def forward(self, features: Tensor, invariant_edge_feats: Tensor, basis: Tensor):
+        with nvtx_range(f'VersatileConvSE3'):
+            num_edges = features.shape[0]
+            in_dim = features.shape[2]
+            with nvtx_range(f'RadialProfile'):
+                radial_weights = self.radial_func(invariant_edge_feats) \
+                    .view(-1, self.channels_out, self.channels_in * self.freq_sum)
+
+            if basis is not None:
+                # This block performs the einsum n i l, n o i f, n l f k -> n o k
+                out_dim = basis.shape[-1]
+                if self.fuse_level != ConvSE3FuseLevel.FULL:
+                    out_dim += out_dim % 2 - 1  # Account for padded basis
+                basis_view = basis.view(num_edges, in_dim, -1)
+                tmp = (features @ basis_view).view(num_edges, -1, basis.shape[-1])
+                return (radial_weights @ tmp)[:, :, :out_dim]
+            else:
+                # k = l = 0 non-fused case
+                return radial_weights @ features
+
+
+class ConvSE3(nn.Module):
+    """
+    SE(3)-equivariant graph convolution (Tensor Field Network convolution).
+    This convolution can map an arbitrary input Fiber to an arbitrary output Fiber, while preserving equivariance.
+    Features of different degrees interact together to produce output features.
+
+    Note 1:
+        The option is given to not pool the output. This means that the convolution sum over neighbors will not be
+        done, and the returned features will be edge features instead of node features.
+
+    Note 2:
+        Unlike the original paper and implementation, this convolution can handle edge feature of degree greater than 0.
+        Input edge features are concatenated with input source node features before the kernel is applied.
+     """
+
+    def __init__(
+            self,
+            fiber_in: Fiber,
+            fiber_out: Fiber,
+            fiber_edge: Fiber,
+            pool: bool = True,
+            use_layer_norm: bool = False,
+            self_interaction: bool = False,
+            max_degree: int = 4,
+            fuse_level: ConvSE3FuseLevel = ConvSE3FuseLevel.FULL,
+            allow_fused_output: bool = False
+    ):
+        """
+        :param fiber_in:           Fiber describing the input features
+        :param fiber_out:          Fiber describing the output features
+        :param fiber_edge:         Fiber describing the edge features (node distances excluded)
+        :param pool:               If True, compute final node features by averaging incoming edge features
+        :param use_layer_norm:     Apply layer normalization between MLP layers
+        :param self_interaction:   Apply self-interaction of nodes
+        :param max_degree:         Maximum degree used in the bases computation
+        :param fuse_level:         Maximum fuse level to use in TFN convolutions
+        :param allow_fused_output: Allow the module to output a fused representation of features
+        """
+        super().__init__()
+        self.pool = pool
+        self.fiber_in = fiber_in
+        self.fiber_out = fiber_out
+        self.self_interaction = self_interaction
+        self.max_degree = max_degree
+        self.allow_fused_output = allow_fused_output
+
+        # channels_in: account for the concatenation of edge features
+        channels_in_set = set([f.channels + fiber_edge[f.degree] * (f.degree > 0) for f in self.fiber_in])
+        channels_out_set = set([f.channels for f in self.fiber_out])
+        unique_channels_in = (len(channels_in_set) == 1)
+        unique_channels_out = (len(channels_out_set) == 1)
+        degrees_up_to_max = list(range(max_degree + 1))
+        common_args = dict(edge_dim=fiber_edge[0] + 1, use_layer_norm=use_layer_norm)
+
+        if fuse_level.value >= ConvSE3FuseLevel.FULL.value and \
+                unique_channels_in and fiber_in.degrees == degrees_up_to_max and \
+                unique_channels_out and fiber_out.degrees == degrees_up_to_max:
+            # Single fused convolution
+            self.used_fuse_level = ConvSE3FuseLevel.FULL
+
+            sum_freq = sum([
+                degree_to_dim(min(d_in, d_out))
+                for d_in, d_out in product(degrees_up_to_max, degrees_up_to_max)
+            ])
+
+            self.conv = VersatileConvSE3(sum_freq, list(channels_in_set)[0], list(channels_out_set)[0],
+                                         fuse_level=self.used_fuse_level, **common_args)
+
+        elif fuse_level.value >= ConvSE3FuseLevel.PARTIAL.value and \
+                unique_channels_in and fiber_in.degrees == degrees_up_to_max:
+            # Convolutions fused per output degree
+            self.used_fuse_level = ConvSE3FuseLevel.PARTIAL
+            self.conv_out = nn.ModuleDict()
+            for d_out, c_out in fiber_out:
+                sum_freq = sum([degree_to_dim(min(d_out, d)) for d in fiber_in.degrees])
+                self.conv_out[str(d_out)] = VersatileConvSE3(sum_freq, list(channels_in_set)[0], c_out,
+                                                             fuse_level=self.used_fuse_level, **common_args)
+
+        elif fuse_level.value >= ConvSE3FuseLevel.PARTIAL.value and \
+                unique_channels_out and fiber_out.degrees == degrees_up_to_max:
+            # Convolutions fused per input degree
+            self.used_fuse_level = ConvSE3FuseLevel.PARTIAL
+            self.conv_in = nn.ModuleDict()
+            for d_in, c_in in fiber_in:
+                sum_freq = sum([degree_to_dim(min(d_in, d)) for d in fiber_out.degrees])
+                self.conv_in[str(d_in)] = VersatileConvSE3(sum_freq, c_in, list(channels_out_set)[0],
+                                                           fuse_level=ConvSE3FuseLevel.FULL, **common_args)
+                                                           #fuse_level=self.used_fuse_level, **common_args)
+        else:
+            # Use pairwise TFN convolutions
+            self.used_fuse_level = ConvSE3FuseLevel.NONE
+            self.conv = nn.ModuleDict()
+            for (degree_in, channels_in), (degree_out, channels_out) in (self.fiber_in * self.fiber_out):
+                dict_key = f'{degree_in},{degree_out}'
+                channels_in_new = channels_in + fiber_edge[degree_in] * (degree_in > 0)
+                sum_freq = degree_to_dim(min(degree_in, degree_out))
+                self.conv[dict_key] = VersatileConvSE3(sum_freq, channels_in_new, channels_out,
+                                                       fuse_level=self.used_fuse_level, **common_args)
+
+        if self_interaction:
+            self.to_kernel_self = nn.ParameterDict()
+            for degree_out, channels_out in fiber_out:
+                if fiber_in[degree_out]:
+                    self.to_kernel_self[str(degree_out)] = nn.Parameter(
+                        torch.randn(channels_out, fiber_in[degree_out]) / np.sqrt(fiber_in[degree_out]))
+
+    def forward(
+            self,
+            node_feats: Dict[str, Tensor],
+            edge_feats: Dict[str, Tensor],
+            graph: DGLGraph,
+            basis: Dict[str, Tensor]
+    ):
+        with nvtx_range(f'ConvSE3'):
+            invariant_edge_feats = edge_feats['0'].squeeze(-1)
+            src, dst = graph.edges()
+            out = {}
+            in_features = []
+
+            # Fetch all input features from edge and node features
+            for degree_in in self.fiber_in.degrees:
+                src_node_features = node_feats[str(degree_in)][src]
+                if degree_in > 0 and str(degree_in) in edge_feats:
+                    # Handle edge features of any type by concatenating them to node features
+                    src_node_features = torch.cat([src_node_features, edge_feats[str(degree_in)]], dim=1)
+                in_features.append(src_node_features)
+
+            if self.used_fuse_level == ConvSE3FuseLevel.FULL:
+                in_features_fused = torch.cat(in_features, dim=-1)
+                out = self.conv(in_features_fused, invariant_edge_feats, basis['fully_fused'])
+
+                if not self.allow_fused_output or self.self_interaction or self.pool:
+                    out = unfuse_features(out, self.fiber_out.degrees)
+
+            elif self.used_fuse_level == ConvSE3FuseLevel.PARTIAL and hasattr(self, 'conv_out'):
+                in_features_fused = torch.cat(in_features, dim=-1)
+                for degree_out in self.fiber_out.degrees:
+                    out[str(degree_out)] = self.conv_out[str(degree_out)](in_features_fused, invariant_edge_feats,
+                                                                          basis[f'out{degree_out}_fused'])
+
+            elif self.used_fuse_level == ConvSE3FuseLevel.PARTIAL and hasattr(self, 'conv_in'):
+                out = 0
+                for degree_in, feature in zip(self.fiber_in.degrees, in_features):
+                    out += self.conv_in[str(degree_in)](feature, invariant_edge_feats,
+                                                        basis[f'in{degree_in}_fused'])
+                if not self.allow_fused_output or self.self_interaction or self.pool:
+                    out = unfuse_features(out, self.fiber_out.degrees)
+            else:
+                # Fallback to pairwise TFN convolutions
+                for degree_out in self.fiber_out.degrees:
+                    out_feature = 0
+                    for degree_in, feature in zip(self.fiber_in.degrees, in_features):
+                        dict_key = f'{degree_in},{degree_out}'
+                        out_feature = out_feature + self.conv[dict_key](feature, invariant_edge_feats,
+                                                                        basis.get(dict_key, None))
+                    out[str(degree_out)] = out_feature
+
+            for degree_out in self.fiber_out.degrees:
+                if self.self_interaction and str(degree_out) in self.to_kernel_self:
+                    with nvtx_range(f'self interaction'):
+                        dst_features = node_feats[str(degree_out)][dst]
+                        kernel_self = self.to_kernel_self[str(degree_out)]
+                        out[str(degree_out)] += kernel_self @ dst_features
+
+                if self.pool:
+                    with nvtx_range(f'pooling'):
+                        if isinstance(out, dict):
+                            out[str(degree_out)] = dgl.ops.copy_e_sum(graph, out[str(degree_out)])
+                        else:
+                            out = dgl.ops.copy_e_sum(graph, out)
+            return out

env/SE3Transformer/build/lib/se3_transformer/model/layers/linear.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from typing import Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from se3_transformer.model.fiber import Fiber
+
+
+class LinearSE3(nn.Module):
+    """
+    Graph Linear SE(3)-equivariant layer, equivalent to a 1x1 convolution.
+    Maps a fiber to a fiber with the same degrees (channels may be different).
+    No interaction between degrees, but interaction between channels.
+
+    type-0 features (C_0 channels) ────> Linear(bias=False) ────> type-0 features (C'_0 channels)
+    type-1 features (C_1 channels) ────> Linear(bias=False) ────> type-1 features (C'_1 channels)
+                                                 :
+    type-k features (C_k channels) ────> Linear(bias=False) ────> type-k features (C'_k channels)
+    """
+
+    def __init__(self, fiber_in: Fiber, fiber_out: Fiber):
+        super().__init__()
+        self.weights = nn.ParameterDict({
+            str(degree_out): nn.Parameter(
+                torch.randn(channels_out, fiber_in[degree_out]) / np.sqrt(fiber_in[degree_out]))
+            for degree_out, channels_out in fiber_out
+        })
+
+    def forward(self, features: Dict[str, Tensor], *args, **kwargs) -> Dict[str, Tensor]:
+        return {
+            degree: self.weights[degree] @ features[degree]
+            for degree, weight in self.weights.items()
+        }

env/SE3Transformer/build/lib/se3_transformer/model/layers/norm.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.model.fiber import Fiber
+
+
+class NormSE3(nn.Module):
+    """
+    Norm-based SE(3)-equivariant nonlinearity.
+
+                 ┌──> feature_norm ──> LayerNorm() ──> ReLU() ──┐
+    feature_in ──┤                                              * ──> feature_out
+                 └──> feature_phase ────────────────────────────┘
+    """
+
+    NORM_CLAMP = 2 ** -24  # Minimum positive subnormal for FP16
+
+    def __init__(self, fiber: Fiber, nonlinearity: nn.Module = nn.ReLU()):
+        super().__init__()
+        self.fiber = fiber
+        self.nonlinearity = nonlinearity
+
+        if len(set(fiber.channels)) == 1:
+            # Fuse all the layer normalizations into a group normalization
+            self.group_norm = nn.GroupNorm(num_groups=len(fiber.degrees), num_channels=sum(fiber.channels))
+        else:
+            # Use multiple layer normalizations
+            self.layer_norms = nn.ModuleDict({
+                str(degree): nn.LayerNorm(channels)
+                for degree, channels in fiber
+            })
+
+    def forward(self, features: Dict[str, Tensor], *args, **kwargs) -> Dict[str, Tensor]:
+        with nvtx_range('NormSE3'):
+            output = {}
+            if hasattr(self, 'group_norm'):
+                # Compute per-degree norms of features
+                norms = [features[str(d)].norm(dim=-1, keepdim=True).clamp(min=self.NORM_CLAMP)
+                         for d in self.fiber.degrees]
+                fused_norms = torch.cat(norms, dim=-2)
+
+                # Transform the norms only
+                new_norms = self.nonlinearity(self.group_norm(fused_norms.squeeze(-1))).unsqueeze(-1)
+                new_norms = torch.chunk(new_norms, chunks=len(self.fiber.degrees), dim=-2)
+
+                # Scale features to the new norms
+                for norm, new_norm, d in zip(norms, new_norms, self.fiber.degrees):
+                    output[str(d)] = features[str(d)] / norm * new_norm
+            else:
+                for degree, feat in features.items():
+                    norm = feat.norm(dim=-1, keepdim=True).clamp(min=self.NORM_CLAMP)
+                    new_norm = self.nonlinearity(self.layer_norms[degree](norm.squeeze(-1)).unsqueeze(-1))
+                    output[degree] = new_norm * feat / norm
+
+            return output

env/SE3Transformer/build/lib/se3_transformer/model/layers/pooling.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from typing import Dict, Literal
+
+import torch.nn as nn
+from dgl import DGLGraph
+from dgl.nn.pytorch import AvgPooling, MaxPooling
+from torch import Tensor
+
+
+class GPooling(nn.Module):
+    """
+    Graph max/average pooling on a given feature type.
+    The average can be taken for any feature type, and equivariance will be maintained.
+    The maximum can only be taken for invariant features (type 0).
+    If you want max-pooling for type > 0 features, look into Vector Neurons.
+    """
+
+    def __init__(self, feat_type: int = 0, pool: Literal['max', 'avg'] = 'max'):
+        """
+        :param feat_type: Feature type to pool
+        :param pool: Type of pooling: max or avg
+        """
+        super().__init__()
+        assert pool in ['max', 'avg'], f'Unknown pooling: {pool}'
+        assert feat_type == 0 or pool == 'avg', 'Max pooling on type > 0 features will break equivariance'
+        self.feat_type = feat_type
+        self.pool = MaxPooling() if pool == 'max' else AvgPooling()
+
+    def forward(self, features: Dict[str, Tensor], graph: DGLGraph, **kwargs) -> Tensor:
+        pooled = self.pool(graph, features[str(self.feat_type)])
+        return pooled.squeeze(dim=-1)

env/SE3Transformer/build/lib/se3_transformer/model/transformer.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import logging
+from typing import Optional, Literal, Dict
+
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from torch import Tensor
+
+from se3_transformer.model.basis import get_basis, update_basis_with_fused
+from se3_transformer.model.layers.attention import AttentionBlockSE3
+from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel
+from se3_transformer.model.layers.norm import NormSE3
+from se3_transformer.model.layers.pooling import GPooling
+from se3_transformer.runtime.utils import str2bool
+from se3_transformer.model.fiber import Fiber
+
+
+class Sequential(nn.Sequential):
+    """ Sequential module with arbitrary forward args and kwargs. Used to pass graph, basis and edge features. """
+
+    def forward(self, input, *args, **kwargs):
+        for module in self:
+            input = module(input, *args, **kwargs)
+        return input
+
+
+def get_populated_edge_features(relative_pos: Tensor, edge_features: Optional[Dict[str, Tensor]] = None):
+    """ Add relative positions to existing edge features """
+    edge_features = edge_features.copy() if edge_features else {}
+    r = relative_pos.norm(dim=-1, keepdim=True)
+    if '0' in edge_features:
+        edge_features['0'] = torch.cat([edge_features['0'], r[..., None]], dim=1)
+    else:
+        edge_features['0'] = r[..., None]
+
+    return edge_features
+
+
+class SE3Transformer(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 fiber_in: Fiber,
+                 fiber_hidden: Fiber,
+                 fiber_out: Fiber,
+                 num_heads: int,
+                 channels_div: int,
+                 fiber_edge: Fiber = Fiber({}),
+                 return_type: Optional[int] = None,
+                 pooling: Optional[Literal['avg', 'max']] = None,
+                 norm: bool = True,
+                 use_layer_norm: bool = True,
+                 tensor_cores: bool = False,
+                 low_memory: bool = False,
+                 **kwargs):
+        """
+        :param num_layers:          Number of attention layers
+        :param fiber_in:            Input fiber description
+        :param fiber_hidden:        Hidden fiber description
+        :param fiber_out:           Output fiber description
+        :param fiber_edge:          Input edge fiber description
+        :param num_heads:           Number of attention heads
+        :param channels_div:        Channels division before feeding to attention layer
+        :param return_type:         Return only features of this type
+        :param pooling:             'avg' or 'max' graph pooling before MLP layers
+        :param norm:                Apply a normalization layer after each attention block
+        :param use_layer_norm:      Apply layer normalization between MLP layers
+        :param tensor_cores:        True if using Tensor Cores (affects the use of fully fused convs, and padded bases)
+        :param low_memory:          If True, will use slower ops that use less memory
+        """
+        super().__init__()
+        self.num_layers = num_layers
+        self.fiber_edge = fiber_edge
+        self.num_heads = num_heads
+        self.channels_div = channels_div
+        self.return_type = return_type
+        self.pooling = pooling
+        self.max_degree = max(*fiber_in.degrees, *fiber_hidden.degrees, *fiber_out.degrees)
+        self.tensor_cores = tensor_cores
+        self.low_memory = low_memory
+
+        if low_memory and not tensor_cores:
+            logging.warning('Low memory mode will have no effect with no Tensor Cores')
+
+        # Fully fused convolutions when using Tensor Cores (and not low memory mode)
+        fuse_level = ConvSE3FuseLevel.FULL if tensor_cores and not low_memory else ConvSE3FuseLevel.PARTIAL
+
+        graph_modules = []
+        for i in range(num_layers):
+            graph_modules.append(AttentionBlockSE3(fiber_in=fiber_in,
+                                                   fiber_out=fiber_hidden,
+                                                   fiber_edge=fiber_edge,
+                                                   num_heads=num_heads,
+                                                   channels_div=channels_div,
+                                                   use_layer_norm=use_layer_norm,
+                                                   max_degree=self.max_degree,
+                                                   fuse_level=fuse_level))
+            if norm:
+                graph_modules.append(NormSE3(fiber_hidden))
+            fiber_in = fiber_hidden
+
+        graph_modules.append(ConvSE3(fiber_in=fiber_in,
+                                     fiber_out=fiber_out,
+                                     fiber_edge=fiber_edge,
+                                     self_interaction=True,
+                                     use_layer_norm=use_layer_norm,
+                                     max_degree=self.max_degree))
+        self.graph_modules = Sequential(*graph_modules)
+
+        if pooling is not None:
+            assert return_type is not None, 'return_type must be specified when pooling'
+            self.pooling_module = GPooling(pool=pooling, feat_type=return_type)
+
+    def forward(self, graph: DGLGraph, node_feats: Dict[str, Tensor],
+                edge_feats: Optional[Dict[str, Tensor]] = None,
+                basis: Optional[Dict[str, Tensor]] = None):
+        # Compute bases in case they weren't precomputed as part of the data loading
+        basis = basis or get_basis(graph.edata['rel_pos'], max_degree=self.max_degree, compute_gradients=False,
+                                   use_pad_trick=self.tensor_cores and not self.low_memory,
+                                   amp=torch.is_autocast_enabled())
+
+        # Add fused bases (per output degree, per input degree, and fully fused) to the dict
+        basis = update_basis_with_fused(basis, self.max_degree, use_pad_trick=self.tensor_cores and not self.low_memory,
+                                        fully_fused=self.tensor_cores and not self.low_memory)
+
+        edge_feats = get_populated_edge_features(graph.edata['rel_pos'], edge_feats)
+
+        node_feats = self.graph_modules(node_feats, edge_feats, graph=graph, basis=basis)
+
+        if self.pooling is not None:
+            return self.pooling_module(node_feats, graph=graph)
+
+        if self.return_type is not None:
+            return node_feats[str(self.return_type)]
+
+        return node_feats
+
+    @staticmethod
+    def add_argparse_args(parser):
+        parser.add_argument('--num_layers', type=int, default=7,
+                            help='Number of stacked Transformer layers')
+        parser.add_argument('--num_heads', type=int, default=8,
+                            help='Number of heads in self-attention')
+        parser.add_argument('--channels_div', type=int, default=2,
+                            help='Channels division before feeding to attention layer')
+        parser.add_argument('--pooling', type=str, default=None, const=None, nargs='?', choices=['max', 'avg'],
+                            help='Type of graph pooling')
+        parser.add_argument('--norm', type=str2bool, nargs='?', const=True, default=False,
+                            help='Apply a normalization layer after each attention block')
+        parser.add_argument('--use_layer_norm', type=str2bool, nargs='?', const=True, default=False,
+                            help='Apply layer normalization between MLP layers')
+        parser.add_argument('--low_memory', type=str2bool, nargs='?', const=True, default=False,
+                            help='If true, will use fused ops that are slower but that use less memory '
+                                 '(expect 25 percent less memory). '
+                                 'Only has an effect if AMP is enabled on Volta GPUs, or if running on Ampere GPUs')
+
+        return parser
+
+
+class SE3TransformerPooled(nn.Module):
+    def __init__(self,
+                 fiber_in: Fiber,
+                 fiber_out: Fiber,
+                 fiber_edge: Fiber,
+                 num_degrees: int,
+                 num_channels: int,
+                 output_dim: int,
+                 **kwargs):
+        super().__init__()
+        kwargs['pooling'] = kwargs['pooling'] or 'max'
+        self.transformer = SE3Transformer(
+            fiber_in=fiber_in,
+            fiber_hidden=Fiber.create(num_degrees, num_channels),
+            fiber_out=fiber_out,
+            fiber_edge=fiber_edge,
+            return_type=0,
+            **kwargs
+        )
+
+        n_out_features = fiber_out.num_features
+        self.mlp = nn.Sequential(
+            nn.Linear(n_out_features, n_out_features),
+            nn.ReLU(),
+            nn.Linear(n_out_features, output_dim)
+        )
+
+    def forward(self, graph, node_feats, edge_feats, basis=None):
+        feats = self.transformer(graph, node_feats, edge_feats, basis).squeeze(-1)
+        y = self.mlp(feats).squeeze(-1)
+        return y
+
+    @staticmethod
+    def add_argparse_args(parent_parser):
+        parser = parent_parser.add_argument_group("Model architecture")
+        SE3Transformer.add_argparse_args(parser)
+        parser.add_argument('--num_degrees',
+                            help='Number of degrees to use. Hidden features will have types [0, ..., num_degrees - 1]',
+                            type=int, default=4)
+        parser.add_argument('--num_channels', help='Number of channels for the hidden features', type=int, default=32)
+        return parent_parser

env/SE3Transformer/build/lib/se3_transformer/runtime/arguments.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import argparse
+import pathlib
+
+from se3_transformer.data_loading import QM9DataModule
+from se3_transformer.model import SE3TransformerPooled
+from se3_transformer.runtime.utils import str2bool
+
+PARSER = argparse.ArgumentParser(description='SE(3)-Transformer')
+
+paths = PARSER.add_argument_group('Paths')
+paths.add_argument('--data_dir', type=pathlib.Path, default=pathlib.Path('./data'),
+                   help='Directory where the data is located or should be downloaded')
+paths.add_argument('--log_dir', type=pathlib.Path, default=pathlib.Path('/results'),
+                   help='Directory where the results logs should be saved')
+paths.add_argument('--dllogger_name', type=str, default='dllogger_results.json',
+                   help='Name for the resulting DLLogger JSON file')
+paths.add_argument('--save_ckpt_path', type=pathlib.Path, default=None,
+                   help='File where the checkpoint should be saved')
+paths.add_argument('--load_ckpt_path', type=pathlib.Path, default=None,
+                   help='File of the checkpoint to be loaded')
+
+optimizer = PARSER.add_argument_group('Optimizer')
+optimizer.add_argument('--optimizer', choices=['adam', 'sgd', 'lamb'], default='adam')
+optimizer.add_argument('--learning_rate', '--lr', dest='learning_rate', type=float, default=0.002)
+optimizer.add_argument('--min_learning_rate', '--min_lr', dest='min_learning_rate', type=float, default=None)
+optimizer.add_argument('--momentum', type=float, default=0.9)
+optimizer.add_argument('--weight_decay', type=float, default=0.1)
+
+PARSER.add_argument('--epochs', type=int, default=100, help='Number of training epochs')
+PARSER.add_argument('--batch_size', type=int, default=240, help='Batch size')
+PARSER.add_argument('--seed', type=int, default=None, help='Set a seed globally')
+PARSER.add_argument('--num_workers', type=int, default=8, help='Number of dataloading workers')
+
+PARSER.add_argument('--amp', type=str2bool, nargs='?', const=True, default=False, help='Use Automatic Mixed Precision')
+PARSER.add_argument('--gradient_clip', type=float, default=None, help='Clipping of the gradient norms')
+PARSER.add_argument('--accumulate_grad_batches', type=int, default=1, help='Gradient accumulation')
+PARSER.add_argument('--ckpt_interval', type=int, default=-1, help='Save a checkpoint every N epochs')
+PARSER.add_argument('--eval_interval', dest='eval_interval', type=int, default=1,
+                    help='Do an evaluation round every N epochs')
+PARSER.add_argument('--silent', type=str2bool, nargs='?', const=True, default=False,
+                    help='Minimize stdout output')
+
+PARSER.add_argument('--benchmark', type=str2bool, nargs='?', const=True, default=False,
+                    help='Benchmark mode')
+
+QM9DataModule.add_argparse_args(PARSER)
+SE3TransformerPooled.add_argparse_args(PARSER)

env/SE3Transformer/build/lib/se3_transformer/runtime/callbacks.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+import torch
+
+from se3_transformer.runtime.loggers import Logger
+from se3_transformer.runtime.metrics import MeanAbsoluteError
+
+
+class BaseCallback(ABC):
+    def on_fit_start(self, optimizer, args):
+        pass
+
+    def on_fit_end(self):
+        pass
+
+    def on_epoch_end(self):
+        pass
+
+    def on_batch_start(self):
+        pass
+
+    def on_validation_step(self, input, target, pred):
+        pass
+
+    def on_validation_end(self, epoch=None):
+        pass
+
+    def on_checkpoint_load(self, checkpoint):
+        pass
+
+    def on_checkpoint_save(self, checkpoint):
+        pass
+
+
+class LRSchedulerCallback(BaseCallback):
+    def __init__(self, logger: Optional[Logger] = None):
+        self.logger = logger
+        self.scheduler = None
+
+    @abstractmethod
+    def get_scheduler(self, optimizer, args):
+        pass
+
+    def on_fit_start(self, optimizer, args):
+        self.scheduler = self.get_scheduler(optimizer, args)
+
+    def on_checkpoint_load(self, checkpoint):
+        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+
+    def on_checkpoint_save(self, checkpoint):
+        checkpoint['scheduler_state_dict'] = self.scheduler.state_dict()
+
+    def on_epoch_end(self):
+        if self.logger is not None:
+            self.logger.log_metrics({'learning rate': self.scheduler.get_last_lr()[0]}, step=self.scheduler.last_epoch)
+        self.scheduler.step()
+
+
+class QM9MetricCallback(BaseCallback):
+    """ Logs the rescaled mean absolute error for QM9 regression tasks """
+
+    def __init__(self, logger, targets_std, prefix=''):
+        self.mae = MeanAbsoluteError()
+        self.logger = logger
+        self.targets_std = targets_std
+        self.prefix = prefix
+        self.best_mae = float('inf')
+
+    def on_validation_step(self, input, target, pred):
+        self.mae(pred.detach(), target.detach())
+
+    def on_validation_end(self, epoch=None):
+        mae = self.mae.compute() * self.targets_std
+        logging.info(f'{self.prefix} MAE: {mae}')
+        self.logger.log_metrics({f'{self.prefix} MAE': mae}, epoch)
+        self.best_mae = min(self.best_mae, mae)
+
+    def on_fit_end(self):
+        if self.best_mae != float('inf'):
+            self.logger.log_metrics({f'{self.prefix} best MAE': self.best_mae})
+
+
+class QM9LRSchedulerCallback(LRSchedulerCallback):
+    def __init__(self, logger, epochs):
+        super().__init__(logger)
+        self.epochs = epochs
+
+    def get_scheduler(self, optimizer, args):
+        min_lr = args.min_learning_rate if args.min_learning_rate else args.learning_rate / 10.0
+        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, self.epochs, eta_min=min_lr)
+
+
+class PerformanceCallback(BaseCallback):
+    def __init__(self, logger, batch_size: int, warmup_epochs: int = 1, mode: str = 'train'):
+        self.batch_size = batch_size
+        self.warmup_epochs = warmup_epochs
+        self.epoch = 0
+        self.timestamps = []
+        self.mode = mode
+        self.logger = logger
+
+    def on_batch_start(self):
+        if self.epoch >= self.warmup_epochs:
+            self.timestamps.append(time.time() * 1000.0)
+
+    def _log_perf(self):
+        stats = self.process_performance_stats()
+        for k, v in stats.items():
+            logging.info(f'performance {k}: {v}')
+
+        self.logger.log_metrics(stats)
+
+    def on_epoch_end(self):
+        self.epoch += 1
+
+    def on_fit_end(self):
+        if self.epoch > self.warmup_epochs:
+            self._log_perf()
+            self.timestamps = []
+
+    def process_performance_stats(self):
+        timestamps = np.asarray(self.timestamps)
+        deltas = np.diff(timestamps)
+        throughput = (self.batch_size / deltas).mean()
+        stats = {
+            f"throughput_{self.mode}": throughput,
+            f"latency_{self.mode}_mean": deltas.mean(),
+            f"total_time_{self.mode}": timestamps[-1] - timestamps[0],
+        }
+        for level in [90, 95, 99]:
+            stats.update({f"latency_{self.mode}_{level}": np.percentile(deltas, level)})
+
+        return stats

env/SE3Transformer/build/lib/se3_transformer/runtime/gpu_affinity.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import collections
+import itertools
+import math
+import os
+import pathlib
+import re
+
+import pynvml
+
+
+class Device:
+    # assumes nvml returns list of 64 bit ints
+    _nvml_affinity_elements = math.ceil(os.cpu_count() / 64)
+
+    def __init__(self, device_idx):
+        super().__init__()
+        self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
+
+    def get_name(self):
+        return pynvml.nvmlDeviceGetName(self.handle)
+
+    def get_uuid(self):
+        return pynvml.nvmlDeviceGetUUID(self.handle)
+
+    def get_cpu_affinity(self):
+        affinity_string = ""
+        for j in pynvml.nvmlDeviceGetCpuAffinity(self.handle, Device._nvml_affinity_elements):
+            # assume nvml returns list of 64 bit ints
+            affinity_string = "{:064b}".format(j) + affinity_string
+
+        affinity_list = [int(x) for x in affinity_string]
+        affinity_list.reverse()  # so core 0 is in 0th element of list
+
+        ret = [i for i, e in enumerate(affinity_list) if e != 0]
+        return ret
+
+
+def get_thread_siblings_list():
+    """
+    Returns a list of 2-element integer tuples representing pairs of
+    hyperthreading cores.
+    """
+    path = "/sys/devices/system/cpu/cpu*/topology/thread_siblings_list"
+    thread_siblings_list = []
+    pattern = re.compile(r"(\d+)\D(\d+)")
+    for fname in pathlib.Path(path[0]).glob(path[1:]):
+        with open(fname) as f:
+            content = f.read().strip()
+            res = pattern.findall(content)
+            if res:
+                pair = tuple(map(int, res[0]))
+                thread_siblings_list.append(pair)
+    return thread_siblings_list
+
+
+def check_socket_affinities(socket_affinities):
+    # sets of cores should be either identical or disjoint
+    for i, j in itertools.product(socket_affinities, socket_affinities):
+        if not set(i) == set(j) and not set(i).isdisjoint(set(j)):
+            raise RuntimeError(f"Sets of cores should be either identical or disjoint, " f"but got {i} and {j}.")
+
+
+def get_socket_affinities(nproc_per_node, exclude_unavailable_cores=True):
+    devices = [Device(i) for i in range(nproc_per_node)]
+    socket_affinities = [dev.get_cpu_affinity() for dev in devices]
+
+    if exclude_unavailable_cores:
+        available_cores = os.sched_getaffinity(0)
+        socket_affinities = [list(set(affinity) & available_cores) for affinity in socket_affinities]
+
+    check_socket_affinities(socket_affinities)
+
+    return socket_affinities
+
+
+def set_socket_affinity(gpu_id):
+    """
+    The process is assigned with all available logical CPU cores from the CPU
+    socket connected to the GPU with a given id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    dev = Device(gpu_id)
+    affinity = dev.get_cpu_affinity()
+    os.sched_setaffinity(0, affinity)
+
+
+def set_single_affinity(gpu_id):
+    """
+    The process is assigned with the first available logical CPU core from the
+    list of all CPU cores from the CPU socket connected to the GPU with a given
+    id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    dev = Device(gpu_id)
+    affinity = dev.get_cpu_affinity()
+
+    # exclude unavailable cores
+    available_cores = os.sched_getaffinity(0)
+    affinity = list(set(affinity) & available_cores)
+    os.sched_setaffinity(0, affinity[:1])
+
+
+def set_single_unique_affinity(gpu_id, nproc_per_node):
+    """
+    The process is assigned with a single unique available physical CPU core
+    from the list of all CPU cores from the CPU socket connected to the GPU with
+    a given id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    socket_affinities = get_socket_affinities(nproc_per_node)
+
+    siblings_list = get_thread_siblings_list()
+    siblings_dict = dict(siblings_list)
+
+    # remove siblings
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities[idx] = list(set(socket_affinity) - set(siblings_dict.values()))
+
+    affinities = []
+    assigned = []
+
+    for socket_affinity in socket_affinities:
+        for core in socket_affinity:
+            if core not in assigned:
+                affinities.append([core])
+                assigned.append(core)
+                break
+    os.sched_setaffinity(0, affinities[gpu_id])
+
+
+def set_socket_unique_affinity(gpu_id, nproc_per_node, mode, balanced=True):
+    """
+    The process is assigned with an unique subset of available physical CPU
+    cores from the CPU socket connected to a GPU with a given id.
+    Assignment automatically includes hyperthreading siblings (if siblings are
+    available).
+
+    Args:
+        gpu_id: index of a GPU
+        nproc_per_node: total number of processes per node
+        mode: mode
+        balanced: assign an equal number of physical cores to each process
+    """
+    socket_affinities = get_socket_affinities(nproc_per_node)
+
+    siblings_list = get_thread_siblings_list()
+    siblings_dict = dict(siblings_list)
+
+    # remove hyperthreading siblings
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities[idx] = list(set(socket_affinity) - set(siblings_dict.values()))
+
+    socket_affinities_to_device_ids = collections.defaultdict(list)
+
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities_to_device_ids[tuple(socket_affinity)].append(idx)
+
+    # compute minimal number of physical cores per GPU across all GPUs and
+    # sockets, code assigns this number of cores per GPU if balanced == True
+    min_physical_cores_per_gpu = min(
+        [len(cores) // len(gpus) for cores, gpus in socket_affinities_to_device_ids.items()]
+    )
+
+    for socket_affinity, device_ids in socket_affinities_to_device_ids.items():
+        devices_per_group = len(device_ids)
+        if balanced:
+            cores_per_device = min_physical_cores_per_gpu
+            socket_affinity = socket_affinity[: devices_per_group * min_physical_cores_per_gpu]
+        else:
+            cores_per_device = len(socket_affinity) // devices_per_group
+
+        for group_id, device_id in enumerate(device_ids):
+            if device_id == gpu_id:
+
+                # In theory there should be no difference in performance between
+                # 'interleaved' and 'continuous' pattern on Intel-based DGX-1,
+                # but 'continuous' should be better for DGX A100 because on AMD
+                # Rome 4 consecutive cores are sharing L3 cache.
+                # TODO: code doesn't attempt to automatically detect layout of
+                # L3 cache, also external environment may already exclude some
+                # cores, this code makes no attempt to detect it and to align
+                # mapping to multiples of 4.
+
+                if mode == "interleaved":
+                    affinity = list(socket_affinity[group_id::devices_per_group])
+                elif mode == "continuous":
+                    affinity = list(socket_affinity[group_id * cores_per_device: (group_id + 1) * cores_per_device])
+                else:
+                    raise RuntimeError("Unknown set_socket_unique_affinity mode")
+
+                # unconditionally reintroduce hyperthreading siblings, this step
+                # may result in a different numbers of logical cores assigned to
+                # each GPU even if balanced == True (if hyperthreading siblings
+                # aren't available for a subset of cores due to some external
+                # constraints, siblings are re-added unconditionally, in the
+                # worst case unavailable logical core will be ignored by
+                # os.sched_setaffinity().
+                affinity += [siblings_dict[aff] for aff in affinity if aff in siblings_dict]
+                os.sched_setaffinity(0, affinity)
+
+
+def set_affinity(gpu_id, nproc_per_node, mode="socket_unique_continuous", balanced=True):
+    """
+    The process is assigned with a proper CPU affinity which matches hardware
+    architecture on a given platform. Usually it improves and stabilizes
+    performance of deep learning training workloads.
+
+    This function assumes that the workload is running in multi-process
+    single-device mode (there are multiple training processes and each process
+    is running on a single GPU), which is typical for multi-GPU training
+    workloads using `torch.nn.parallel.DistributedDataParallel`.
+
+    Available affinity modes:
+    * 'socket' - the process is assigned with all available logical CPU cores
+    from the CPU socket connected to the GPU with a given id.
+    * 'single' - the process is assigned with the first available logical CPU
+    core from the list of all CPU cores from the CPU socket connected to the GPU
+    with a given id (multiple GPUs could be assigned with the same CPU core).
+    * 'single_unique' - the process is assigned with a single unique available
+    physical CPU core from the list of all CPU cores from the CPU socket
+    connected to the GPU with a given id.
+    * 'socket_unique_interleaved' - the process is assigned with an unique
+    subset of available physical CPU cores from the CPU socket connected to a
+    GPU with a given id, hyperthreading siblings are included automatically,
+    cores are assigned with interleaved indexing pattern
+    * 'socket_unique_continuous' - (the default) the process is assigned with an
+    unique subset of available physical CPU cores from the CPU socket connected
+    to a GPU with a given id, hyperthreading siblings are included
+    automatically, cores are assigned with continuous indexing pattern
+
+    'socket_unique_continuous' is the recommended mode for deep learning
+    training workloads on NVIDIA DGX machines.
+
+    Args:
+        gpu_id: integer index of a GPU
+        nproc_per_node: number of processes per node
+        mode: affinity mode
+        balanced: assign an equal number of physical cores to each process,
+            affects only 'socket_unique_interleaved' and
+            'socket_unique_continuous' affinity modes
+
+    Returns a set of logical CPU cores on which the process is eligible to run.
+
+    Example:
+
+    import argparse
+    import os
+
+    import gpu_affinity
+    import torch
+
+
+    def main():
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            '--local_rank',
+            type=int,
+            default=os.getenv('LOCAL_RANK', 0),
+        )
+        args = parser.parse_args()
+
+        nproc_per_node = torch.cuda.device_count()
+
+        affinity = gpu_affinity.set_affinity(args.local_rank, nproc_per_node)
+        print(f'{args.local_rank}: core affinity: {affinity}')
+
+
+    if __name__ == "__main__":
+        main()
+
+    Launch the example with:
+    python -m torch.distributed.launch --nproc_per_node <#GPUs> example.py
+
+
+    WARNING: On DGX A100 only a half of CPU cores have direct access to GPUs.
+    This function restricts execution only to the CPU cores directly connected
+    to GPUs, so on DGX A100 it will limit the code to half of CPU cores and half
+    of CPU memory bandwidth (which may be fine for many DL models).
+    """
+    pynvml.nvmlInit()
+
+    if mode == "socket":
+        set_socket_affinity(gpu_id)
+    elif mode == "single":
+        set_single_affinity(gpu_id)
+    elif mode == "single_unique":
+        set_single_unique_affinity(gpu_id, nproc_per_node)
+    elif mode == "socket_unique_interleaved":
+        set_socket_unique_affinity(gpu_id, nproc_per_node, "interleaved", balanced)
+    elif mode == "socket_unique_continuous":
+        set_socket_unique_affinity(gpu_id, nproc_per_node, "continuous", balanced)
+    else:
+        raise RuntimeError("Unknown affinity mode")
+
+    affinity = os.sched_getaffinity(0)
+    return affinity

env/SE3Transformer/build/lib/se3_transformer/runtime/inference.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from typing import List
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from se3_transformer.runtime import gpu_affinity
+from se3_transformer.runtime.arguments import PARSER
+from se3_transformer.runtime.callbacks import BaseCallback
+from se3_transformer.runtime.loggers import DLLogger
+from se3_transformer.runtime.utils import to_cuda, get_local_rank
+
+
+@torch.inference_mode()
+def evaluate(model: nn.Module,
+             dataloader: DataLoader,
+             callbacks: List[BaseCallback],
+             args):
+    model.eval()
+    for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), unit='batch', desc=f'Evaluation',
+                         leave=False, disable=(args.silent or get_local_rank() != 0)):
+        *input, target = to_cuda(batch)
+
+        for callback in callbacks:
+            callback.on_batch_start()
+
+        with torch.cuda.amp.autocast(enabled=args.amp):
+            pred = model(*input)
+
+            for callback in callbacks:
+                callback.on_validation_step(input, target, pred)
+
+
+if __name__ == '__main__':
+    from se3_transformer.runtime.callbacks import QM9MetricCallback, PerformanceCallback
+    from se3_transformer.runtime.utils import init_distributed, seed_everything
+    from se3_transformer.model import SE3TransformerPooled, Fiber
+    from se3_transformer.data_loading import QM9DataModule
+    import torch.distributed as dist
+    import logging
+    import sys
+
+    is_distributed = init_distributed()
+    local_rank = get_local_rank()
+    args = PARSER.parse_args()
+
+    logging.getLogger().setLevel(logging.CRITICAL if local_rank != 0 or args.silent else logging.INFO)
+
+    logging.info('====== SE(3)-Transformer ======')
+    logging.info('|  Inference on the test set  |')
+    logging.info('===============================')
+
+    if not args.benchmark and args.load_ckpt_path is None:
+        logging.error('No load_ckpt_path provided, you need to provide a saved model to evaluate')
+        sys.exit(1)
+
+    if args.benchmark:
+        logging.info('Running benchmark mode with one warmup pass')
+
+    if args.seed is not None:
+        seed_everything(args.seed)
+
+    major_cc, minor_cc = torch.cuda.get_device_capability()
+
+    logger = DLLogger(args.log_dir, filename=args.dllogger_name)
+    datamodule = QM9DataModule(**vars(args))
+    model = SE3TransformerPooled(
+        fiber_in=Fiber({0: datamodule.NODE_FEATURE_DIM}),
+        fiber_out=Fiber({0: args.num_degrees * args.num_channels}),
+        fiber_edge=Fiber({0: datamodule.EDGE_FEATURE_DIM}),
+        output_dim=1,
+        tensor_cores=(args.amp and major_cc >= 7) or major_cc >= 8,  # use Tensor Cores more effectively
+        **vars(args)
+    )
+    callbacks = [QM9MetricCallback(logger, targets_std=datamodule.targets_std, prefix='test')]
+
+    model.to(device=torch.cuda.current_device())
+    if args.load_ckpt_path is not None:
+        checkpoint = torch.load(str(args.load_ckpt_path), map_location={'cuda:0': f'cuda:{local_rank}'})
+        model.load_state_dict(checkpoint['state_dict'])
+
+    if is_distributed:
+        nproc_per_node = torch.cuda.device_count()
+        affinity = gpu_affinity.set_affinity(local_rank, nproc_per_node)
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+
+    test_dataloader = datamodule.test_dataloader() if not args.benchmark else datamodule.train_dataloader()
+    evaluate(model,
+             test_dataloader,
+             callbacks,
+             args)
+
+    for callback in callbacks:
+        callback.on_validation_end()
+
+    if args.benchmark:
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        callbacks = [PerformanceCallback(logger, args.batch_size * world_size, warmup_epochs=1, mode='inference')]
+        for _ in range(6):
+            evaluate(model,
+                     test_dataloader,
+                     callbacks,
+                     args)
+            callbacks[0].on_epoch_end()
+
+        callbacks[0].on_fit_end()

env/SE3Transformer/build/lib/se3_transformer/runtime/loggers.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import pathlib
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, Any, Callable, Optional
+
+import dllogger
+import torch.distributed as dist
+import wandb
+from dllogger import Verbosity
+
+from se3_transformer.runtime.utils import rank_zero_only
+
+
+class Logger(ABC):
+    @rank_zero_only
+    @abstractmethod
+    def log_hyperparams(self, params):
+        pass
+
+    @rank_zero_only
+    @abstractmethod
+    def log_metrics(self, metrics, step=None):
+        pass
+
+    @staticmethod
+    def _sanitize_params(params):
+        def _sanitize(val):
+            if isinstance(val, Callable):
+                try:
+                    _val = val()
+                    if isinstance(_val, Callable):
+                        return val.__name__
+                    return _val
+                except Exception:
+                    return getattr(val, "__name__", None)
+            elif isinstance(val, pathlib.Path) or isinstance(val, Enum):
+                return str(val)
+            return val
+
+        return {key: _sanitize(val) for key, val in params.items()}
+
+
+class LoggerCollection(Logger):
+    def __init__(self, loggers):
+        super().__init__()
+        self.loggers = loggers
+
+    def __getitem__(self, index):
+        return [logger for logger in self.loggers][index]
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        for logger in self.loggers:
+            logger.log_metrics(metrics, step)
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        for logger in self.loggers:
+            logger.log_hyperparams(params)
+
+
+class DLLogger(Logger):
+    def __init__(self, save_dir: pathlib.Path, filename: str):
+        super().__init__()
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            save_dir.mkdir(parents=True, exist_ok=True)
+            dllogger.init(
+                backends=[dllogger.JSONStreamBackend(Verbosity.DEFAULT, str(save_dir / filename))])
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        params = self._sanitize_params(params)
+        dllogger.log(step="PARAMETER", data=params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        if step is None:
+            step = tuple()
+
+        dllogger.log(step=step, data=metrics)
+
+
+class WandbLogger(Logger):
+    def __init__(
+            self,
+            name: str,
+            save_dir: pathlib.Path,
+            id: Optional[str] = None,
+            project: Optional[str] = None
+    ):
+        super().__init__()
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            save_dir.mkdir(parents=True, exist_ok=True)
+            self.experiment = wandb.init(name=name,
+                                         project=project,
+                                         id=id,
+                                         dir=str(save_dir),
+                                         resume='allow',
+                                         anonymous='must')
+
+    @rank_zero_only
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        params = self._sanitize_params(params)
+        self.experiment.config.update(params, allow_val_change=True)
+
+    @rank_zero_only
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        if step is not None:
+            self.experiment.log({**metrics, 'epoch': step})
+        else:
+            self.experiment.log(metrics)

env/SE3Transformer/build/lib/se3_transformer/runtime/metrics.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+
+class Metric(ABC):
+    """ Metric class with synchronization capabilities similar to TorchMetrics """
+
+    def __init__(self):
+        self.states = {}
+
+    def add_state(self, name: str, default: Tensor):
+        assert name not in self.states
+        self.states[name] = default.clone()
+        setattr(self, name, default)
+
+    def synchronize(self):
+        if dist.is_initialized():
+            for state in self.states:
+                dist.all_reduce(getattr(self, state), op=dist.ReduceOp.SUM, group=dist.group.WORLD)
+
+    def __call__(self, *args, **kwargs):
+        self.update(*args, **kwargs)
+
+    def reset(self):
+        for name, default in self.states.items():
+            setattr(self, name, default.clone())
+
+    def compute(self):
+        self.synchronize()
+        value = self._compute().item()
+        self.reset()
+        return value
+
+    @abstractmethod
+    def _compute(self):
+        pass
+
+    @abstractmethod
+    def update(self, preds: Tensor, targets: Tensor):
+        pass
+
+
+class MeanAbsoluteError(Metric):
+    def __init__(self):
+        super().__init__()
+        self.add_state('error', torch.tensor(0, dtype=torch.float32, device='cuda'))
+        self.add_state('total', torch.tensor(0, dtype=torch.int32, device='cuda'))
+
+    def update(self, preds: Tensor, targets: Tensor):
+        preds = preds.detach()
+        n = preds.shape[0]
+        error = torch.abs(preds.view(n, -1) - targets.view(n, -1)).sum()
+        self.total += n
+        self.error += error
+
+    def _compute(self):
+        return self.error / self.total

env/SE3Transformer/build/lib/se3_transformer/runtime/training.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import logging
+import pathlib
+from typing import List
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from apex.optimizers import FusedAdam, FusedLAMB
+from torch.nn.modules.loss import _Loss
+from torch.nn.parallel import DistributedDataParallel
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader, DistributedSampler
+from tqdm import tqdm
+
+from se3_transformer.data_loading import QM9DataModule
+from se3_transformer.model import SE3TransformerPooled
+from se3_transformer.model.fiber import Fiber
+from se3_transformer.runtime import gpu_affinity
+from se3_transformer.runtime.arguments import PARSER
+from se3_transformer.runtime.callbacks import QM9MetricCallback, QM9LRSchedulerCallback, BaseCallback, \
+    PerformanceCallback
+from se3_transformer.runtime.inference import evaluate
+from se3_transformer.runtime.loggers import LoggerCollection, DLLogger, WandbLogger, Logger
+from se3_transformer.runtime.utils import to_cuda, get_local_rank, init_distributed, seed_everything, \
+    using_tensor_cores, increase_l2_fetch_granularity
+
+
+def save_state(model: nn.Module, optimizer: Optimizer, epoch: int, path: pathlib.Path, callbacks: List[BaseCallback]):
+    """ Saves model, optimizer and epoch states to path (only once per node) """
+    if get_local_rank() == 0:
+        state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
+        checkpoint = {
+            'state_dict': state_dict,
+            'optimizer_state_dict': optimizer.state_dict(),
+            'epoch': epoch
+        }
+        for callback in callbacks:
+            callback.on_checkpoint_save(checkpoint)
+
+        torch.save(checkpoint, str(path))
+        logging.info(f'Saved checkpoint to {str(path)}')
+
+
+def load_state(model: nn.Module, optimizer: Optimizer, path: pathlib.Path, callbacks: List[BaseCallback]):
+    """ Loads model, optimizer and epoch states from path """
+    checkpoint = torch.load(str(path), map_location={'cuda:0': f'cuda:{get_local_rank()}'})
+    if isinstance(model, DistributedDataParallel):
+        model.module.load_state_dict(checkpoint['state_dict'])
+    else:
+        model.load_state_dict(checkpoint['state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+    for callback in callbacks:
+        callback.on_checkpoint_load(checkpoint)
+
+    logging.info(f'Loaded checkpoint from {str(path)}')
+    return checkpoint['epoch']
+
+
+def train_epoch(model, train_dataloader, loss_fn, epoch_idx, grad_scaler, optimizer, local_rank, callbacks, args):
+    losses = []
+    for i, batch in tqdm(enumerate(train_dataloader), total=len(train_dataloader), unit='batch',
+                         desc=f'Epoch {epoch_idx}', disable=(args.silent or local_rank != 0)):
+        *inputs, target = to_cuda(batch)
+
+        for callback in callbacks:
+            callback.on_batch_start()
+
+        with torch.cuda.amp.autocast(enabled=args.amp):
+            pred = model(*inputs)
+            loss = loss_fn(pred, target) / args.accumulate_grad_batches
+
+        grad_scaler.scale(loss).backward()
+
+        # gradient accumulation
+        if (i + 1) % args.accumulate_grad_batches == 0 or (i + 1) == len(train_dataloader):
+            if args.gradient_clip:
+                grad_scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.gradient_clip)
+
+            grad_scaler.step(optimizer)
+            grad_scaler.update()
+            optimizer.zero_grad()
+
+        losses.append(loss.item())
+
+    return np.mean(losses)
+
+
+def train(model: nn.Module,
+          loss_fn: _Loss,
+          train_dataloader: DataLoader,
+          val_dataloader: DataLoader,
+          callbacks: List[BaseCallback],
+          logger: Logger,
+          args):
+    device = torch.cuda.current_device()
+    model.to(device=device)
+    local_rank = get_local_rank()
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+
+    model.train()
+    grad_scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
+    if args.optimizer == 'adam':
+        optimizer = FusedAdam(model.parameters(), lr=args.learning_rate, betas=(args.momentum, 0.999),
+                              weight_decay=args.weight_decay)
+    elif args.optimizer == 'lamb':
+        optimizer = FusedLAMB(model.parameters(), lr=args.learning_rate, betas=(args.momentum, 0.999),
+                              weight_decay=args.weight_decay)
+    else:
+        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum,
+                                    weight_decay=args.weight_decay)
+
+    epoch_start = load_state(model, optimizer, args.load_ckpt_path, callbacks) if args.load_ckpt_path else 0
+
+    for callback in callbacks:
+        callback.on_fit_start(optimizer, args)
+
+    for epoch_idx in range(epoch_start, args.epochs):
+        if isinstance(train_dataloader.sampler, DistributedSampler):
+            train_dataloader.sampler.set_epoch(epoch_idx)
+
+        loss = train_epoch(model, train_dataloader, loss_fn, epoch_idx, grad_scaler, optimizer, local_rank, callbacks, args)
+        if dist.is_initialized():
+            loss = torch.tensor(loss, dtype=torch.float, device=device)
+            torch.distributed.all_reduce(loss)
+            loss = (loss / world_size).item()
+
+        logging.info(f'Train loss: {loss}')
+        logger.log_metrics({'train loss': loss}, epoch_idx)
+
+        for callback in callbacks:
+            callback.on_epoch_end()
+
+        if not args.benchmark and args.save_ckpt_path is not None and args.ckpt_interval > 0 \
+                and (epoch_idx + 1) % args.ckpt_interval == 0:
+            save_state(model, optimizer, epoch_idx, args.save_ckpt_path, callbacks)
+
+        if not args.benchmark and args.eval_interval > 0 and (epoch_idx + 1) % args.eval_interval == 0:
+            evaluate(model, val_dataloader, callbacks, args)
+            model.train()
+
+            for callback in callbacks:
+                callback.on_validation_end(epoch_idx)
+
+    if args.save_ckpt_path is not None and not args.benchmark:
+        save_state(model, optimizer, args.epochs, args.save_ckpt_path, callbacks)
+
+    for callback in callbacks:
+        callback.on_fit_end()
+
+
+def print_parameters_count(model):
+    num_params_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    logging.info(f'Number of trainable parameters: {num_params_trainable}')
+
+
+if __name__ == '__main__':
+    is_distributed = init_distributed()
+    local_rank = get_local_rank()
+    args = PARSER.parse_args()
+
+    logging.getLogger().setLevel(logging.CRITICAL if local_rank != 0 or args.silent else logging.INFO)
+
+    logging.info('====== SE(3)-Transformer ======')
+    logging.info('|      Training procedure     |')
+    logging.info('===============================')
+
+    if args.seed is not None:
+        logging.info(f'Using seed {args.seed}')
+        seed_everything(args.seed)
+
+    logger = LoggerCollection([
+        DLLogger(save_dir=args.log_dir, filename=args.dllogger_name),
+        WandbLogger(name=f'QM9({args.task})', save_dir=args.log_dir, project='se3-transformer')
+    ])
+
+    datamodule = QM9DataModule(**vars(args))
+    model = SE3TransformerPooled(
+        fiber_in=Fiber({0: datamodule.NODE_FEATURE_DIM}),
+        fiber_out=Fiber({0: args.num_degrees * args.num_channels}),
+        fiber_edge=Fiber({0: datamodule.EDGE_FEATURE_DIM}),
+        output_dim=1,
+        tensor_cores=using_tensor_cores(args.amp),  # use Tensor Cores more effectively
+        **vars(args)
+    )
+    loss_fn = nn.L1Loss()
+
+    if args.benchmark:
+        logging.info('Running benchmark mode')
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        callbacks = [PerformanceCallback(logger, args.batch_size * world_size)]
+    else:
+        callbacks = [QM9MetricCallback(logger, targets_std=datamodule.targets_std, prefix='validation'),
+                     QM9LRSchedulerCallback(logger, epochs=args.epochs)]
+
+    if is_distributed:
+        gpu_affinity.set_affinity(gpu_id=get_local_rank(), nproc_per_node=torch.cuda.device_count())
+
+    print_parameters_count(model)
+    logger.log_hyperparams(vars(args))
+    increase_l2_fetch_granularity()
+    train(model,
+          loss_fn,
+          datamodule.train_dataloader(),
+          datamodule.val_dataloader(),
+          callbacks,
+          logger,
+          args)
+
+    logging.info('Training finished successfully')

env/SE3Transformer/build/lib/se3_transformer/runtime/utils.py

@@ -0,0 +1,130 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import argparse
+import ctypes
+import logging
+import os
+import random
+from functools import wraps
+from typing import Union, List, Dict
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+
+def aggregate_residual(feats1, feats2, method: str):
+    """ Add or concatenate two fiber features together. If degrees don't match, will use the ones of feats2. """
+    if method in ['add', 'sum']:
+        return {k: (v + feats1[k]) if k in feats1 else v for k, v in feats2.items()}
+    elif method in ['cat', 'concat']:
+        return {k: torch.cat([v, feats1[k]], dim=1) if k in feats1 else v for k, v in feats2.items()}
+    else:
+        raise ValueError('Method must be add/sum or cat/concat')
+
+
+def degree_to_dim(degree: int) -> int:
+    return 2 * degree + 1
+
+
+def unfuse_features(features: Tensor, degrees: List[int]) -> Dict[str, Tensor]:
+    return dict(zip(map(str, degrees), features.split([degree_to_dim(deg) for deg in degrees], dim=-1)))
+
+
+def str2bool(v: Union[bool, str]) -> bool:
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def to_cuda(x):
+    """ Try to convert a Tensor, a collection of Tensors or a DGLGraph to CUDA """
+    if isinstance(x, Tensor):
+        return x.cuda(non_blocking=True)
+    elif isinstance(x, tuple):
+        return (to_cuda(v) for v in x)
+    elif isinstance(x, list):
+        return [to_cuda(v) for v in x]
+    elif isinstance(x, dict):
+        return {k: to_cuda(v) for k, v in x.items()}
+    else:
+        # DGLGraph or other objects
+        return x.to(device=torch.cuda.current_device())
+
+
+def get_local_rank() -> int:
+    return int(os.environ.get('LOCAL_RANK', 0))
+
+
+def init_distributed() -> bool:
+    world_size = int(os.environ.get('WORLD_SIZE', 1))
+    distributed = world_size > 1
+    if distributed:
+        backend = 'nccl' if torch.cuda.is_available() else 'gloo'
+        dist.init_process_group(backend=backend, init_method='env://')
+        if backend == 'nccl':
+            torch.cuda.set_device(get_local_rank())
+        else:
+            logging.warning('Running on CPU only!')
+        assert torch.distributed.is_initialized()
+    return distributed
+
+
+def increase_l2_fetch_granularity():
+    # maximum fetch granularity of L2: 128 bytes
+    _libcudart = ctypes.CDLL('libcudart.so')
+    # set device limit on the current device
+    # cudaLimitMaxL2FetchGranularity = 0x05
+    pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
+    _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128))
+    _libcudart.cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
+    assert pValue.contents.value == 128
+
+
+def seed_everything(seed):
+    seed = int(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def rank_zero_only(fn):
+    @wraps(fn)
+    def wrapped_fn(*args, **kwargs):
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            return fn(*args, **kwargs)
+
+    return wrapped_fn
+
+
+def using_tensor_cores(amp: bool) -> bool:
+    major_cc, minor_cc = torch.cuda.get_device_capability()
+    return (amp and major_cc >= 7) or major_cc >= 8

env/SE3Transformer/build/lib/tests/test_equivariance.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import torch
+
+from se3_transformer.model import SE3Transformer
+from se3_transformer.model.fiber import Fiber
+from tests.utils import get_random_graph, assign_relative_pos, get_max_diff, rot
+
+# Tolerances for equivariance error abs( f(x) @ R  -  f(x @ R) )
+TOL = 1e-3
+CHANNELS, NODES = 32, 512
+
+
+def _get_outputs(model, R):
+    feats0 = torch.randn(NODES, CHANNELS, 1)
+    feats1 = torch.randn(NODES, CHANNELS, 3)
+
+    coords = torch.randn(NODES, 3)
+    graph = get_random_graph(NODES)
+    if torch.cuda.is_available():
+        feats0 = feats0.cuda()
+        feats1 = feats1.cuda()
+        R = R.cuda()
+        coords = coords.cuda()
+        graph = graph.to('cuda')
+        model.cuda()
+
+    graph1 = assign_relative_pos(graph, coords)
+    out1 = model(graph1, {'0': feats0, '1': feats1}, {})
+    graph2 = assign_relative_pos(graph, coords @ R)
+    out2 = model(graph2, {'0': feats0, '1': feats1 @ R}, {})
+
+    return out1, out2
+
+
+def _get_model(**kwargs):
+    return SE3Transformer(
+        num_layers=4,
+        fiber_in=Fiber.create(2, CHANNELS),
+        fiber_hidden=Fiber.create(3, CHANNELS),
+        fiber_out=Fiber.create(2, CHANNELS),
+        fiber_edge=Fiber({}),
+        num_heads=8,
+        channels_div=2,
+        **kwargs
+    )
+
+
+def test_equivariance():
+    model = _get_model()
+    R = rot(*torch.rand(3))
+    if torch.cuda.is_available():
+        R = R.cuda()
+    out1, out2 = _get_outputs(model, R)
+
+    assert torch.allclose(out2['0'], out1['0'], atol=TOL), \
+        f'type-0 features should be invariant {get_max_diff(out1["0"], out2["0"])}'
+    assert torch.allclose(out2['1'], (out1['1'] @ R), atol=TOL), \
+        f'type-1 features should be equivariant {get_max_diff(out1["1"] @ R, out2["1"])}'
+
+
+def test_equivariance_pooled():
+    model = _get_model(pooling='avg', return_type=1)
+    R = rot(*torch.rand(3))
+    if torch.cuda.is_available():
+        R = R.cuda()
+    out1, out2 = _get_outputs(model, R)
+
+    assert torch.allclose(out2, (out1 @ R), atol=TOL), \
+        f'type-1 features should be equivariant {get_max_diff(out1 @ R, out2)}'
+
+
+def test_invariance_pooled():
+    model = _get_model(pooling='avg', return_type=0)
+    R = rot(*torch.rand(3))
+    if torch.cuda.is_available():
+        R = R.cuda()
+    out1, out2 = _get_outputs(model, R)
+
+    assert torch.allclose(out2, out1, atol=TOL), \
+        f'type-0 features should be invariant {get_max_diff(out1, out2)}'

env/SE3Transformer/build/lib/tests/utils.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import dgl
+import torch
+
+
+def get_random_graph(N, num_edges_factor=18):
+    graph = dgl.transform.remove_self_loop(dgl.rand_graph(N, N * num_edges_factor))
+    return graph
+
+
+def assign_relative_pos(graph, coords):
+    src, dst = graph.edges()
+    graph.edata['rel_pos'] = coords[src] - coords[dst]
+    return graph
+
+
+def get_max_diff(a, b):
+    return (a - b).abs().max().item()
+
+
+def rot_z(gamma):
+    return torch.tensor([
+        [torch.cos(gamma), -torch.sin(gamma), 0],
+        [torch.sin(gamma), torch.cos(gamma), 0],
+        [0, 0, 1]
+    ], dtype=gamma.dtype)
+
+
+def rot_y(beta):
+    return torch.tensor([
+        [torch.cos(beta), 0, torch.sin(beta)],
+        [0, 1, 0],
+        [-torch.sin(beta), 0, torch.cos(beta)]
+    ], dtype=beta.dtype)
+
+
+def rot(alpha, beta, gamma):
+    return rot_z(alpha) @ rot_y(beta) @ rot_z(gamma)