Function _device_agnostic_dispatch

tests/testing_utils.py:1491–1502

(device: str, dispatch_table: dict[str, Callable], *args, **kwargs)

# This dispatches a defined function according to the accelerator from the function definitions.
def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs):
    if device not in dispatch_table:
        return dispatch_table["default"](*args, **kwargs)

    fn = dispatch_table[device]

    # Some device agnostic functions return values. Need to guard against 'None' instead at
    # user level
    if not callable(fn):
        return fn

    return fn(*args, **kwargs)


# These are callables which automatically dispatch the function specific to the accelerator