Switch function/class implementation between Numpy & PyTorch:?

user_4570472 0
Viewed 76

I have a function (actually a class, but for simplicity, let's pretend it's a function) that uses several NumPy operations that all exist in PyTorch e.g. np.add and I also want a PyTorch version of the function. I'm trying to avoid duplicating my code, so I want to know:

Is there a way for me to dynamically switch a function's execution back and forth between NumPy and PyTorch without needing duplicate implementations?

For a toy example, suppose my function is:

def foo_numpy(x: np.ndarray, y: np.ndarray) -> np.ndarray:
  return np.add(x, y)

I could define a PyTorch equivalent:

def foo_torch(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
  return torch.add(x, y)

Could I somehow define a function like:

def foo(x, y, mode: str = 'numpy'):
  if mode == 'numpy':
    return np.add(x, y)
  elif mode == 'torch':
    return torch.add(x, y)
  else:
    raise ValueError

without needing the if-else statement?

Edit: what about something like the following?

def foo(x, y, mode: str = 'numpy'):
  if mode == 'numpy':
    lib = np
  elif mode == 'torch':
    lib = torch
  else:
    raise ValueError
  return lib.add(x, y)
pytorch numpy python
1 Answers

Instead of using a string, you can use a boolean (bool) value to represent the mode you want to use i.e. False (0) representing NumPy and True (1) representing PyTorch. One can then use ternary operators to further shrink the if statements.

def foo(x, y, mode: bool = 0):
    lib = torch if mode else np
    return lib.add(x, y)

If you want to switch back and forth between the two in a class you can do something similar

class Example:

    def __init__(self):
        self._mode = True
    
    def switchMode(self):
        self._mode = !self._mode

    def foo(self, x, y):
        lib = torch if self._mode else np
        return lib.add(x, y)
user_122012430
Related