Remove representation dependencies and refactor codebase

[ziofil]

User description

This PR removes the representations module and refactors the entire codebase to eliminate representation dependencies, simplifying the architecture and improving maintainability. Key changes include: removal of representations.py module, refactoring of all state and transformation classes, fixing JAX tracer leaks, and updating all tests to work with the new structure.

---

PR Type

Enhancement

---

Description

• Remove representations module and refactor architecture
• Update all states and transformations to use ansatz/wires directly
• Fix JAX tracer leaks in vacuum state initialization
• Simplify circuit component structure and imports

---

Changes walkthrough 📝

	Relevant files
Enhancement	17 files circuit_components.py Refactor CircuitComponent to use ansatz/wires directly      +148/-83 base.py Remove representation dependencies from transformation base classes +22/-32  ket.py Refactor Ket class to remove representation dependencies  +30/-9    gaussian_state.py Update gaussian states to use ansatz/wires directly            +25/-28  dm.py Refactor DM class to remove representation dependencies    +3/-4      bsgate.py Update BSgate to use ansatz/wires directly                              +8/-11    dgate.py Update Dgate to use ansatz/wires directly                                +6/-9      displaced_squeezed.py Update DisplacedSqueezed to use ansatz/wires directly        +22/-15  phasenoise.py Update PhaseNoise to use ansatz/wires directly                      +10/-7    mzgate.py Update MZgate to use ansatz/wires directly                              +10/-12  b_to_q.py Update BtoQ to use ansatz/wires directly                                  +12/-13  quadrature_eigenstate.py Update QuadratureEigenstate to use ansatz/wires directly  +14/-11  coherent.py Update Coherent state to use ansatz/wires directly              +17/-11  trace_out.py Update TraceOut to use ansatz/wires directly                          +3/-6      two_mode_squeezed_vacuum.py Update TwoModeSqueezedVacuum to use ansatz/wires directly +13/-10  s2gate.py Update S2gate to use ansatz/wires directly                              +13/-11  squeezed_vacuum.py Update SqueezedVacuum to use ansatz/wires directly              +14/-11
Tests	3 files test_circuit_components.py Update tests for new CircuitComponent structure                    +22/-29  test_transformations_base.py Update transformation tests for new structure                        +10/-15  test_representations.py Update representation tests to use CircuitComponent            +8/-8
Additional files	29 files representations.rst +0/-6      b_to_ch.py +15/-12  b_to_ps.py +14/-10  bargmann_eigenstate.py +10/-7    number.py +5/-5      sauron.py +8/-9      thermal.py +8/-7      vacuum.py +6/-5      amplifier.py +12/-10  attenuator.py +15/-12  cxgate.py +17/-12  czgate.py +16/-11  fockdamping.py +5/-7      gaussrandnoise.py +11/-8    ggate.py +16/-9    identity.py +9/-5      interferometer.py +8/-10    pgate.py +14/-11  realinterferometer.py +13/-11  rgate.py +5/-7      sgate.py +18/-11  representations.py +0/-224  triples.py +4/-3      test_circuits.py +1/-1      test_dm.py +5/-4      test_ket.py +6/-4      test_number.py +1/-1      test_quadrature_eigenstate.py +1/-1      test_wires.py +2/-2

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-merge-pro]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 5 🔵🔵🔵🔵🔵

🏅 Score: 75

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Complex Logic The contract method contains complex logic with multiple conditional paths and error handling that should be carefully validated for correctness, especially the ansatz type checking and batch string operations. return self * other if type(self.ansatz) is type(other.ansatz): self_ansatz = self.ansatz other_ansatz = other.ansatz self_wires = self.wires other_wires = other.wires else: self_fock = self.to_fock() other_fock = other.to_fock() self_ansatz = self_fock.ansatz other_ansatz = other_fock.ansatz self_wires = self_fock.wires other_wires = other_fock.wires # Contract the ansatzes if type(self_ansatz) is not type(other_ansatz): raise ValueError( f"Cannot contract ansatz of type {type(self_ansatz)} with ansatz ", f"of type {type(other_ansatz)}. Please call either `to_fock` or `to_bargmann` ", "on one of the representations.", ) wires_result, _ = self_wires @ other_wires core1, core2, core_out = self_wires.contracted_labels(other_wires) if mode == "zip": from mrmustard.physics.utils import zip_batch_strings eins_str = zip_batch_strings( self_ansatz.batch_dims - self_ansatz._lin_sup, other_ansatz.batch_dims - other_ansatz._lin_sup, ) elif mode == "kron": from mrmustard.physics.utils import outer_product_batch_str eins_str = outer_product_batch_str( self_ansatz.batch_dims - self_ansatz._lin_sup, other_ansatz.batch_dims - other_ansatz._lin_sup, ) batch12, batch_out = eins_str.split("->") batch1, batch2 = batch12.split(",") ansatz = self_ansatz.contract( other_ansatz, list(batch1) + core1, list(batch2) + core2, list(batch_out) + core_out ) return CircuitComponent(ansatz, wires_result) Commented Code Large block of commented-out code in the dm method suggests incomplete refactoring or uncertainty about the implementation approach. # Contract the ansatzes # from mrmustard.physics.utils import zip_batch_strings # self_ansatz = self.ansatz # other_ansatz = self.adjoint.ansatz # self_wires = self.wires # other_wires = self.adjoint.wires # wires_result, _ = self_wires @ other_wires # core1, core2, core_out = self_wires.contracted_labels(other_wires) # eins_str = zip_batch_strings( # self_ansatz.batch_dims - self_ansatz._lin_sup, # other_ansatz.batch_dims - other_ansatz._lin_sup, # ) # batch12, batch_out = eins_str.split("->") # batch1, batch2 = batch12.split(",") # ansatz = self_ansatz.contract( # other_ansatz, list(batch1) + core1, list(batch2) + core2, list(batch_out) + core_out # ) # ret = DM(ansatz, wires_result, name=self.name) # ret.manual_shape = self.manual_shape + self.manual_shape # return ret Debug Code Debug print statement left in production code that should be removed before merging. # print(shape) shape = shape or self.auto_shape()

[qodo-merge-pro]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Handle None ansatz case The method should also handle the case where self.ansatz is None, which would raise AttributeError before reaching the triple attribute access. mrmustard/lab/circuit_components.py [500-503] +if self.ansatz is None: + raise AttributeError("No Bargmann data for this component.") try: return self.ansatz.triple except AttributeError as e: raise AttributeError("No Bargmann data for this component.") from e Apply / Chat Suggestion importance[1-10]: 7 __ Why: The suggestion correctly points out that if self.ansatz is None, accessing self.ansatz.triple will raise an AttributeError. The existing try...except block already catches this AttributeError and raises a new one with a more informative message. The suggested change adds an explicit if self.ansatz is None: check, which makes the handling of the None case more explicit and slightly improves readability, but the existing code is functionally correct.	Medium
General	Improve None handling logic The conditional check should handle the case where self.ansatz is None more explicitly to avoid potential issues with method chaining when ansatz is None. mrmustard/lab/circuit_components.py [140] -ansatz = self.ansatz.reorder(kets + bras).conj if self.ansatz else None +if self.ansatz is None: + ansatz = None +else: + ansatz = self.ansatz.reorder(kets + bras).conj Apply / Chat Suggestion importance[1-10]: 2 __ Why: The suggestion proposes rewriting a ternary expression into a more verbose if/else block. While functionally identical, the existing code is a standard and concise Python idiom for conditional assignment. The suggested change does not improve correctness or fix any bug, and the impact on readability is subjective.	Low
Update

[codecov]

Codecov Report

Attention: Patch coverage is 99.27711% with 3 lines in your changes missing coverage. Please review.

Project coverage is 91.50%. Comparing base (fa2df8c) to head (884daf3).
Report is 1 commits behind head on develop.

Files with missing lines	Patch %	Lines
mrmustard/lab/circuit_components.py	97.27%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop     #605      +/-   ##
===========================================
- Coverage    91.51%   91.50%   -0.02%     
===========================================
  Files          102      101       -1     
  Lines         5837     5863      +26     
===========================================
+ Hits          5342     5365      +23     
- Misses         495      498       +3     

Files with missing lines	Coverage Δ
mrmustard/__init__.py	100.00% <100.00%> (ø)
mrmustard/lab/circuit_components_utils/__init__.py	100.00% <100.00%> (ø)
mrmustard/lab/circuit_components_utils/b_to_ch.py	100.00% <100.00%> (ø)
mrmustard/lab/circuit_components_utils/b_to_ps.py	100.00% <100.00%> (ø)
mrmustard/lab/circuit_components_utils/b_to_q.py	100.00% <100.00%> (ø)
...rmustard/lab/circuit_components_utils/trace_out.py	100.00% <100.00%> (ø)
mrmustard/lab/circuits.py	91.93% <100.00%> (ø)
mrmustard/lab/samplers.py	98.90% <100.00%> (ø)
mrmustard/lab/states/__init__.py	100.00% <100.00%> (ø)
mrmustard/lab/states/bargmann_eigenstate.py	100.00% <100.00%> (ø)
... and 65 more

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fa2df8c...884daf3. Read the comment docs.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[heltluke]

To make it easier to focus on the changes relevant to this PR, it could be nice for all of the isort changes to be handled elsewhere (perhaps in a preliminary PR).

[apchytr]

The .rst file here should be removed

[apchytr]

Why is 312 being removed here?

[apchytr]

Now that there isn't a Representation class I think these tests should be ported to it's respective file

[apchytr]

What was the issue here?

[ziofil]

jax had a tracer error, this way it went away