Comprehensive Benchmark for PyTorch backend

[refraction-ray]

Description

With the PyTorch backend now supporting jit (see issue #4), it is crucial to establish a comprehensive benchmark suite. This suite will be used to evaluate and compare the performance of the PyTorch backend agaisnt TensorFlow, and JAX backends across different hardware environments.

The goal is to measure and analyze both compilation and execution times for a variety of representative workloads, as these can vary significantly depending on the problem size, ML backend, and hardware (CPU vs. GPU).

Tasks

1. Benchmark Suite Development

• Develop a suite of benchmarks for the following components:
  • Circuit
  • DMCircuit
  • MPScircuit
  • AD
  • VMAP
  • all with jit
• Ensure benchmarks are sufficiently large to yield meaningful and stable performance measurements, cotengra contractor is often required for consistent performance.
• Parameterize benchmarks to allow for scaling of problem size.

2. Backend and Hardware Testing Matrix

• Configure the benchmark runner to execute across the following configurations:
  • Backends:
    • PyTorch
    • TensorFlow
    • JAX
  • Hardware:
    • CPU
    • GPU

3. Reporting and Analysis

• Log and report both compilation time and execution time for each test case.
• Analyze the results to identify performance bottlenecks and areas for optimization for each backend.

Acceptance Criteria

• A set of benchmark scripts are created and added potentially to the repository.
• The benchmarks can be executed against all three backends (PyTorch, TensorFlow, JAX) and both devices.
• A clear and reproducible report of the results is generated, comparing the performance of the different backends.

Additional Context

This benchmark suite will be a critical tool for ongoing performance monitoring and will help guide future development and optimization efforts. It will also provide valuable information to users, helping them choose the best backend for their specific use cases and hardware.