import random import math import inspect import itertools import logging from typing import Iterable, Optional, Callable import warnings import numpy as np import torch from e3nn import o3 from e3nn.util.jit import ( compile, get_tracing_inputs, get_compile_mode, _MAKE_TRACING_INPUTS, get_optimization_defaults, set_optimization_defaults, ) from ._argtools import _get_args_in, _get_io_irreps, _transform, _rand_args # pylint: disable=unused-variable # Make a logger for reporting error statistics logger = logging.getLogger(__name__) def _logging_name(func) -> str: """Get a decent string representation of ``func`` for logging""" if inspect.isfunction(func): return func.__name__ else: return repr(func) # The default float tolerance FLOAT_TOLERANCE = {t: torch.as_tensor(v, dtype=t) for t, v in {torch.float32: 1e-3, torch.float64: 1e-9}.items()} try: # If pytest is available, define an e3nn pytest plugin # See https://docs.pytest.org/en/stable/fixture.html#using-fixtures-from-other-projects import pytest @pytest.fixture(scope="session", autouse=True, params=["float32", "float64"]) def float_tolerance(request): """Run all tests with various PyTorch default dtypes. This is a session-wide, autouse fixture — you only need to request it explicitly if a test needs to know the tolerance for the current default dtype. Returns -------- A precision threshold to use for closeness tests. """ old_dtype = torch.get_default_dtype() dtype = {"float32": torch.float32, "float64": torch.float64}[request.param] torch.set_default_dtype(dtype) yield FLOAT_TOLERANCE[dtype] torch.set_default_dtype(old_dtype) except ImportError: pass def random_irreps( n: int = 1, lmax: int = 4, mul_min: int = 0, mul_max: int = 5, len_min: int = 0, len_max: int = 4, clean: bool = False, allow_empty: bool = True, ): r"""Generate random irreps parameters for testing. Parameters ---------- n : int, optional How many to generate; defaults to 1. lmax : int, optional The maximum L to generate (inclusive); defaults to 4. mul_min : int, optional The smallest multiplicity to generate, defaults to 0. mul_max : int, optional The largest multiplicity to generate, defaults to 5. len_min : int, optional The smallest number of irreps to generate, defaults to 0. len_max : int, optional The largest number of irreps to generate, defaults to 4. clean : bool, optional If ``True``, only ``o3.Irreps`` objects will be returned. If ``False`` (the default), ``e3nn.o3.Irreps``-like objects like strings and lists of tuples can be returned. allow_empty : bool, optional Whether to allow generating empty ``e3nn.o3.Irreps``. Returns ------- An irreps-like object if ``n == 1`` or a list of them if ``n > 1`` """ assert n >= 1 assert lmax >= 0 assert mul_min >= 0 assert mul_max >= mul_min if not allow_empty and len_min == 0: len_min = 1 assert len_min >= 0 assert len_max >= len_min out = [] for _ in range(n): this_irreps = [] for _ in range(random.randint(len_min, len_max)): this_irreps.append((random.randint(mul_min, mul_max), (random.randint(0, lmax), random.choice((1, -1))))) if not allow_empty and all(m == 0 for m, _ in this_irreps): this_irreps[-1] = (random.randint(1, mul_max), this_irreps[-1][1]) this_irreps = o3.Irreps(this_irreps) if clean: outtype = "irreps" else: outtype = random.choice(("irreps", "str", "list")) if outtype == "irreps": out.append(this_irreps) elif outtype == "str": out.append(repr(this_irreps)) elif outtype == "list": out.append([(mul_ir.mul, (mul_ir.ir.l, mul_ir.ir.p)) for mul_ir in this_irreps]) if n == 1: return out[0] else: return out def format_equivariance_error(errors: dict) -> str: """Format the dictionary returned by ``equivariance_error`` into a readable string. Parameters ---------- errors : dict A dictionary of errors returned by ``equivariance_error``. Returns ------- A string. """ return "\n".join( "(parity_k={:d}, did_translate={}) -> max error={:.3e} in argument {}".format( int(k[0]), bool(k[1]), float(v.max()), int(v.argmax()) ) for k, v in errors.items() ) def assert_equivariant(func, args_in=None, irreps_in=None, irreps_out=None, tolerance=None, **kwargs) -> dict: r"""Assert that ``func`` is equivariant. Parameters ---------- args_in : list or None the original input arguments for the function. If ``None`` and the function has ``irreps_in`` consisting only of ``o3.Irreps`` and ``'cartesian'``, random test inputs will be generated. irreps_in : object see ``equivariance_error`` irreps_out : object see ``equivariance_error`` tolerance : float or None the threshold below which the equivariance error must fall. If ``None``, (the default), ``FLOAT_TOLERANCE[torch.get_default_dtype()]`` is used. **kwargs : kwargs passed through to ``equivariance_error``. Returns ------- The same as ``equivariance_error``: a dictionary mapping tuples ``(parity_k, did_translate)`` to errors """ # Prevent pytest from showing this function in the traceback __tracebackhide__ = True args_in, irreps_in, irreps_out = _get_args_in(func, args_in=args_in, irreps_in=irreps_in, irreps_out=irreps_out) # Get error errors = equivariance_error(func, args_in=args_in, irreps_in=irreps_in, irreps_out=irreps_out, **kwargs) logger.info( "Tested equivariance of `%s` -- max componentwise errors:\n%s", _logging_name(func), format_equivariance_error(errors), ) # Check it if tolerance is None: tolerance = FLOAT_TOLERANCE[torch.get_default_dtype()] problems = {case: err for case, err in errors.items() if err.max() > tolerance} if len(problems) != 0: errstr = "Largest componentwise equivariance error was too large for: " errstr += format_equivariance_error(problems) assert len(problems) == 0, errstr return errors def equivariance_error( func, args_in, irreps_in=None, irreps_out=None, ntrials: int = 1, do_parity: bool = True, do_translation: bool = True, transform_dtype=torch.float64, ): r"""Get the maximum equivariance error for ``func`` over ``ntrials`` Each trial randomizes the equivariant transformation tested. Parameters ---------- func : callable the function to test args_in : list the original inputs to pass to ``func``. irreps_in : list of `e3nn.o3.Irreps` or `e3nn.o3.Irreps` the input irreps for each of the arguments in ``args_in``. If left as the default of ``None``, ``get_io_irreps`` will be used to try to infer them. If a sequence is provided, valid elements are also the string ``'cartesian'``, which denotes that the corresponding input should be dealt with as cartesian points in 3D, and ``None``, which indicates that the argument should not be transformed. irreps_out : list of `e3nn.o3.Irreps` or `e3nn.o3.Irreps` the out irreps for each of the return values of ``func``. Accepts similar values to ``irreps_in``. ntrials : int run this many trials with random transforms do_parity : bool whether to test parity do_translation : bool whether to test translation for ``'cartesian'`` inputs Returns ------- dictionary mapping tuples ``(parity_k, did_translate)`` to an array of errors, each entry the biggest over all trials for that output, in order. """ irreps_in, irreps_out = _get_io_irreps(func, irreps_in=irreps_in, irreps_out=irreps_out) if do_parity: parity_ks = [0, 1] else: parity_ks = [0] if "cartesian_points" not in irreps_in: # There's nothing to translate do_translation = False if do_translation: do_translation = [False, True] else: do_translation = [False] tests = list(itertools.product(parity_ks, do_translation)) neg_inf = -float("Inf") device = next(t.device for t in args_in if isinstance(t, torch.Tensor)) biggest_errs = {test: torch.full((len(irreps_out),), neg_inf, dtype=transform_dtype, device=device) for test in tests} for trial in range(ntrials): for this_test in tests: parity_k, this_do_translate = this_test # Build a rotation matrix for point data rot_mat = o3.rand_matrix(dtype=transform_dtype) # add parity rot_mat *= (-1) ** parity_k # build translation translation = 10 * torch.randn(1, 3, dtype=rot_mat.dtype) if this_do_translate else 0.0 # Evaluate the function on rotated arguments: rot_args = _transform(args_in, irreps_in, rot_mat, translation) x1 = func(*rot_args) if isinstance(x1, torch.Tensor): x1 = [x1] elif isinstance(x1, (list, tuple)): x1 = list(x1) else: raise TypeError(f"equivariance_error cannot handle output type {type(x1)}") # if `func` was a model, the outputs might be attached in the autograd graph # convert into the transform dtype for computing the difference x1 = [t.detach().to(transform_dtype) for t in x1] # Evaluate the function on the arguments, then apply group action: x2 = func(*args_in) if isinstance(x2, torch.Tensor): x2 = [x2] elif isinstance(x2, (list, tuple)): x2 = list(x2) else: raise TypeError(f"equivariance_error cannot handle output type {type(x2)}") x2 = [t.detach() for t in x2] # confirm sanity assert len(x1) == len(x2) assert len(x1) == len(irreps_out) # apply the group action to x2 # get this in the transform dtype x2 = _transform(x2, irreps_out, rot_mat, translation, output_transform_dtype=True) # compute errors in the transform dtype, # then convert back to default later errors = torch.stack([(a - b).abs().max() for a, b in zip(x1, x2)]) biggest_errs[this_test] = torch.where(errors > biggest_errs[this_test], errors, biggest_errs[this_test]) # convert errors back to default dtype to return: return {k: v.to(torch.get_default_dtype()) for k, v in biggest_errs.items()} # TODO: this is only for things marked with @compile_mode. # Make something else for general script/traceability def assert_auto_jitable( func, error_on_warnings: bool = True, n_trace_checks: int = 2, strict_shapes: bool = True, ): r"""Assert that submodule ``func`` is automatically JITable. Parameters ---------- func : Callable The function to trace. error_on_warnings : bool If True (default), TracerWarnings emitted by ``torch.jit.trace`` will be treated as errors. n_random_tests : int If ``args_in`` is ``None`` and arguments are being automatically generated, this many random arguments will be generated as test inputs for ``torch.jit.trace``. strict_shapes : bool Test that the traced function errors on inputs with feature dimensions that don't match the input irreps. Returns ------- The traced TorchScript function. """ # Prevent pytest from showing this function in the traceback __tracebackhide__ = True if get_compile_mode(func) is None: raise ValueError("assert_auto_jitable is only for modules marked with @compile_mode") # Test tracing with warnings.catch_warnings(): if error_on_warnings: warnings.filterwarnings("error", category=torch.jit.TracerWarning) func_jit = compile(func, n_trace_checks=n_trace_checks) # Confirm that it rejects incorrect shapes # This check only makes sense if all inputs are Tensors with irreps; otherwise we can't know how to modify the arguments # or that our modifications make them wrong. if strict_shapes and not hasattr(func, _MAKE_TRACING_INPUTS): try: all_bad_args = get_tracing_inputs(func, n=1)[0] except ValueError: # couldn't infer, don't check pass else: for method, bad_args in all_bad_args.items(): # Since _rand_args is OK, they're all Irreps style args where changing the feature dimension is wrong bad_which = random.randint(0, len(bad_args) - 1) bad_args = list(bad_args) bad_args[bad_which] = bad_args[bad_which][..., : -random.randint(1, 3)] # make bad shape try: if method == "forward": func_jit(*bad_args) else: getattr(func_jit, method)(*bad_args) except (torch.jit.Error, RuntimeError): # type: ignore # As far as I can tell, there's no good way to introspect TorchScript exceptions. pass else: raise AssertionError("Traced function didn't error on bad input shape") return func_jit def assert_torch_compile( compile_mode: str, func: Callable, *args, **kwargs, ) -> None: r"""Assert that func is torch.compile(fullgraph=True) Parameters ---------- func: Callable thats a functools.partial(torch.nn.Module) *args: func's forward arguments **kwargs: func's forward positional arguments """ # Turning off the torch.jit.script in CodeGenMix to enable torch.compile. jit_mode_before = get_optimization_defaults()["jit_mode"] try: set_optimization_defaults(jit_mode=compile_mode) m = func() torch._dynamo.reset() # Clear cache from previous runs m_pt2 = torch.compile(m, fullgraph=True) m_pt2(*args, **kwargs) finally: set_optimization_defaults(jit_mode=jit_mode_before) return m_pt2 # TODO: custom in_vars, out_vars support def assert_normalized( func: torch.nn.Module, irreps_in=None, irreps_out=None, normalization: str = "component", n_input: int = 10_000, n_weight: Optional[int] = None, weights: Optional[Iterable[torch.nn.Parameter]] = None, atol: float = 0.1, ) -> None: r"""Assert that ``func`` is normalized. See https://docs.e3nn.org/en/stable/guide/normalization.html for more information on the normalization scheme. ``atol``, ``n_input``, and ``n_weight`` may need to be significantly higher in order to converge the statistics to pass the test. Parameters ---------- func : torch.nn.Module the module to test irreps_in : object see ``equivariance_error`` irreps_out : object see ``equivariance_error`` normalization : str, default "component" one of "component" or "norm". Note that this is defined for both the inputs and the outputs; if you need seperate normalizations for input and output please file a feature request. n_input : int, default 10_000 the number of input samples to use for each weight init n_weight : int, default 20 the number of weight initializations to sample weights : optional iterable of parameters the weights to reinitialize ``n_weight`` times. If ``None`` (default), ``func.parameters()`` will be used. atol : float, default 0.1 tolerance for checking moments. Higher values for this prevent explosive computational costs for this test. """ # Prevent pytest from showing this function in the traceback __tracebackhide__ = True if normalization not in ("component", "norm"): raise ValueError(f"invalid normalization `{normalization}`") irreps_in, irreps_out = _get_io_irreps(func, irreps_in=irreps_in, irreps_out=irreps_out) if all(i.num_irreps == 0 for i in irreps_in) or all(i.num_irreps == 0 for i in irreps_out): # Short-circut return if weights is None: if isinstance(func, torch.nn.Module): weights = func.parameters() else: weights = [] weights = list(weights) if len(weights) == 0: assert n_weight is None or n_weight == 1, "Without weights to re-init, n_weight must be 1 or None" n_weight = 1 else: n_weight = 20 if n_weight is None else n_weight with torch.no_grad(): expected_squares = [torch.zeros(irreps.dim) for irreps in irreps_out] n_samples = 0 for weight_init in range(n_weight): # generate weight sample for param in weights: param.normal_() # generate input sample args_in = _rand_args(irreps_in, batch_size=n_input) # args_in gives component normalized irreps if normalization == "norm": for i, irreps in enumerate(irreps_in): for mul_ir, ir_slice in zip(irreps, irreps.slices()): args_in[i][:, ir_slice].div_(math.sqrt(mul_ir.ir.dim)) # run func this_outs = func(*args_in) if not isinstance(this_outs, list) or isinstance(this_outs, tuple): this_outs = (this_outs,) assert len(this_outs) == len(irreps_out) # square this_outs = [e.square() for e in this_outs] # update running average for i in range(len(irreps_out)): assert this_outs[i].shape[0] == n_input update = this_outs[i].sum(dim=0) - n_input * expected_squares[i] update.div_(n_input + n_samples) expected_squares[i].add_(update) n_samples += n_input # check them for expected_square, irreps in zip(expected_squares, irreps_out): if irreps == "cartesian_points" or irreps is None: continue if normalization == "component": targets = [1.0] * len(irreps) elif normalization == "norm": targets = [1.0 / math.sqrt(ir.dim) for _, ir in irreps] for i, (target, ir_slice) in enumerate(zip(targets, irreps.slices())): if ir_slice.start == ir_slice.stop: continue max_componentwise = (expected_square[ir_slice] - target).abs().max().item() logger.info("Tested normalization of %r: max componentwise error %.6f", _logging_name(func), max_componentwise) assert max_componentwise <= atol, ( f"< x_i^2 > !~= {target:.6f} for output irrep #{i}, {irreps[i]}." f"Max componentwise error: {max_componentwise:.6f}" ) def set_random_seeds() -> None: """Set the random seeds to try to get some reproducibility""" torch.manual_seed(0) random.seed(0) np.random.seed(0)