feat: automatic tensor backend and type checks

[yxlao]

Design Goals

Goals for Multi-Backend Support

1. All functions in camtools should be able to process and return both NumPy arrays and Torch tensors. This is achieved automatically by decorators.
2. Torch should NOT be a compulsory dependency of camtools. The library can detect if Torch is installed and enable the corresponding backends.
3. By default, the backend is set to NumPy. You can set the default backend to Torch by using set_backend().

Goals for Type Checking

1. Functions in camtools should be type-hinted with the expected tensor shape and dtype.
2. The shape and dtypes are automatically checked and enforced by decorators. This will replace most of the functions in ct.sanity.

Changes in This PR

• New Dependencies: Introduces ivy, jaxtyping, and einops.
  • ivy: Interoperates between NumPy and Torch. Used internally only.
  • jaxtyping: Provides type hinting for the library.
• Torch as an Optional Dependency: Torch remains optional.
• Backend Management:
  • Use ct.backend.get_backend() and ct.backend.set_backend() to get and set the default backend for camtools. This change is global and affects all subsequent invocations of camtools functions.
  • @ct.backend.tensor_auto_backend tries to derive backend from input arguments. It also converts list inputs to native tensor format if the type hint is a tensor. See test/test_backend.py for examples.
  • @ct.backend.tensor_numpy_backend and @ct.backend.tensor_torch_backend wrap functions to enforce using NumPy and Torch backends by converting input tensors to NumPy or Torch, respectively.
• Shape and Dtype Checking:
  • @ct.backend.tensor_type_check enforces shape and dtype checks based on the type hints for tensors. See test/test_backend.py for examples. This check will eventually replace most of the functions in ct.sanity.

Known Issues

Per our benchmark, the overhead of type-checking and backend conversion is small. However ivy introduces significant overheads. We may look for alternative methods, such as lightweight wrappers to wrap numpy/torch to replace ivy in the future.