{"library":"opt-einsum-fx","title":"Einsum optimization using opt_einsum and PyTorch FX","description":"opt-einsum-fx is a Python library that leverages opt_einsum and PyTorch FX to optimize Einstein summation (einsum) expressions within PyTorch computation graphs. It aims to reduce the overall execution time and memory footprint of complex tensor contractions by intelligently reordering operations. The current version is 0.1.4, with the last release in November 2021, indicating a maintenance-level release cadence.","language":"python","status":"maintenance","last_verified":"Sun May 17","install":{"commands":["pip install opt_einsum_fx"],"cli":null},"imports":["import opt_einsum_fx","from opt_einsum_fx import optimize_einsums_full"],"auth":{"required":false,"env_vars":[]},"quickstart":{"code":"import torch\nimport torch.fx\nimport opt_einsum_fx\n\ndef einmatvecmul(a, b, vec):\n    \"\"\"Batched matrix-matrix-vector product using einsum\"\"\"\n    return torch.einsum(\"zij,zjk,zk->zi\", a, b, vec)\n\n# 1. Create an FX graph module from the function\ngraph_mod = torch.fx.symbolic_trace(einmatvecmul)\n\n# 2. Define example inputs for shape propagation and optimization\n# These shapes are used to determine the optimal contraction path.\nexample_inputs = (\n    torch.randn(7, 4, 5),\n    torch.randn(7, 5, 3),\n    torch.randn(7, 3)\n)\n\n# 3. Optimize the einsums within the FX graph\ngraph_opt = opt_einsum_fx.optimize_einsums_full(\n    model=graph_mod,\n    example_inputs=example_inputs\n)\n\n# 4. (Optional) Print the optimized code to see the changes\nprint(\"Original code:\\n\", graph_mod.code)\nprint(\"Optimized code:\\n\", graph_opt.code)\n\n# 5. Run the optimized graph and verify correctness\noutput_original = graph_mod(*example_inputs)\noutput_optimized = graph_opt(*example_inputs)\n\nassert torch.allclose(output_original, output_optimized)\nprint(\"\\nOptimization successful and outputs match!\")","lang":"python","description":"This quickstart demonstrates how to use `opt_einsum_fx` to optimize a PyTorch function containing an `einsum` operation. It involves symbolic tracing the function with `torch.fx.symbolic_trace`, providing example inputs for shape inference, and then applying `opt_einsum_fx.optimize_einsums_full` to get an optimized graph module. The outputs of the original and optimized graphs are compared to ensure correctness.","tag":null,"tag_description":null,"last_tested":null,"results":[]},"compatibility":{"tag":null,"tag_description":null,"last_tested":"2026-05-17","installed_version":"0.1.4","pypi_latest":"0.1.4","is_stale":false,"summary":{"python_range":"3.10–3.9","success_rate":50,"avg_install_s":74.3,"avg_import_s":4.99,"wheel_type":"wheel"},"results":[{"runtime":"python:3.10-alpine","python_version":"3.10","os_libc":"alpine (musl)","variant":"opt_einsum_fx","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.10-slim","python_version":"3.10","os_libc":"slim (glibc)","variant":"opt_einsum_fx","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"noisy","install_time_s":73.5,"import_time_s":3.72,"mem_mb":59.8,"disk_size":"4.6G"},{"runtime":"python:3.11-alpine","python_version":"3.11","os_libc":"alpine (musl)","variant":"opt_einsum_fx","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.11-slim","python_version":"3.11","os_libc":"slim (glibc)","variant":"opt_einsum_fx","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"noisy","install_time_s":66.4,"import_time_s":5.96,"mem_mb":64.6,"disk_size":"4.7G"},{"runtime":"python:3.12-alpine","python_version":"3.12","os_libc":"alpine (musl)","variant":"opt_einsum_fx","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.12-slim","python_version":"3.12","os_libc":"slim (glibc)","variant":"opt_einsum_fx","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"noisy","install_time_s":62.7,"import_time_s":6.04,"mem_mb":63.5,"disk_size":"4.7G"},{"runtime":"python:3.13-alpine","python_version":"3.13","os_libc":"alpine (musl)","variant":"opt_einsum_fx","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.13-slim","python_version":"3.13","os_libc":"slim (glibc)","variant":"opt_einsum_fx","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"noisy","install_time_s":56.7,"import_time_s":5.18,"mem_mb":64,"disk_size":"4.7G"},{"runtime":"python:3.9-alpine","python_version":"3.9","os_libc":"alpine (musl)","variant":"opt_einsum_fx","exit_code":1,"wheel_type":null,"failure_reason":"build_error","import_side_effects":null,"install_time_s":null,"import_time_s":null,"mem_mb":null,"disk_size":null},{"runtime":"python:3.9-slim","python_version":"3.9","os_libc":"slim (glibc)","variant":"opt_einsum_fx","exit_code":0,"wheel_type":"wheel","failure_reason":null,"import_side_effects":"noisy","install_time_s":112.1,"import_time_s":4.03,"mem_mb":52.3,"disk_size":"6.4G"}]}}