diff --git a/pyproject.toml b/pyproject.toml
index 721ff5124..df77e463b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "ffsim"
 requires-python = ">=3.9"
-version = "0.0.42.dev"
+version = "0.0.44.dev"
 description = "Faster simulations of fermionic quantum circuits."
 readme = "README.md"
 license = { file = "LICENSE" }
diff --git a/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py b/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py
index b2ab966a8..e571928c9 100644
--- a/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py
+++ b/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py
@@ -15,6 +15,7 @@
 
 import numpy as np
 import scipy.linalg
+from pyscf.fci.direct_uhf import make_hdiag
 from scipy.sparse.linalg import LinearOperator
 
 from ffsim._lib import FermionOperator
@@ -174,3 +175,23 @@ def from_fermion_operator(op: FermionOperator) -> DiagonalCoulombHamiltonian:
             diag_coulomb_mats=diag_coulomb_mats,
             constant=constant,
         )
+
+    def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+        """Return the diagonal entries of the Hamiltonian."""
+        if np.iscomplexobj(self.one_body_tensor):
+            raise NotImplementedError(
+                "Computing diagonal of complex diagonal Coulomb Hamiltonian is not yet "
+                "supported."
+            )
+        two_body_tensor_aa = np.zeros((self.norb, self.norb, self.norb, self.norb))
+        two_body_tensor_ab = np.zeros((self.norb, self.norb, self.norb, self.norb))
+        diag_coulomb_mat_aa, diag_coulomb_mat_ab = self.diag_coulomb_mats
+        for p, q in itertools.product(range(self.norb), repeat=2):
+            two_body_tensor_aa[p, p, q, q] = diag_coulomb_mat_aa[p, q]
+            two_body_tensor_ab[p, p, q, q] = diag_coulomb_mat_ab[p, q]
+        one_body_tensor = self.one_body_tensor + 0.5 * np.einsum(
+            "prqr", two_body_tensor_aa
+        )
+        h1e = (one_body_tensor, one_body_tensor)
+        h2e = (two_body_tensor_aa, two_body_tensor_ab, two_body_tensor_aa)
+        return make_hdiag(h1e, h2e, norb=norb, nelec=nelec) + self.constant
diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py
index d8ea415a7..9a65584f9 100644
--- a/python/ffsim/hamiltonians/molecular_hamiltonian.py
+++ b/python/ffsim/hamiltonians/molecular_hamiltonian.py
@@ -158,8 +158,9 @@ def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray:
         if np.iscomplexobj(self.two_body_tensor) or np.iscomplexobj(
             self.one_body_tensor
         ):
-            raise ValueError(
-                "Computing diagonal of complex molecular Hamiltonian is not supported."
+            raise NotImplementedError(
+                "Computing diagonal of complex molecular Hamiltonian is not yet "
+                "supported."
             )
         return (
             make_hdiag(self.one_body_tensor, self.two_body_tensor, norb, nelec)
diff --git a/python/ffsim/molecular_data.py b/python/ffsim/molecular_data.py
index 9664fd6e3..5a30aa051 100644
--- a/python/ffsim/molecular_data.py
+++ b/python/ffsim/molecular_data.py
@@ -26,6 +26,7 @@
 import pyscf
 import pyscf.cc
 import pyscf.ci
+import pyscf.fci
 import pyscf.mcscf
 import pyscf.mp
 import pyscf.symm
@@ -68,6 +69,9 @@ class MolecularData:
             CCSD t2 amplitudes.
         cisd_energy (float | None): The CISD energy.
         cisd_vec (np.ndarray | None): The CISD state vector.
+        sci_energy (float | None): The SCI energy.
+        sci_vec (tuple[np.ndarray, np.ndarray, np.ndarray] | None): The SCI state
+            vector coefficients, spin alpha strings, and spin beta strings.
         fci_energy (float | None): The FCI energy.
         fci_vec (np.ndarray | None): The FCI state vector.
         dipole_integrals (np.ndarray | None): The dipole integrals.
@@ -104,6 +108,9 @@ class MolecularData:
     # CISD data
     cisd_energy: float | None = None
     cisd_vec: np.ndarray | None = None
+    # SCI data
+    sci_energy: float | None = None
+    sci_vec: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None
     # FCI data
     fci_energy: float | None = None
     fci_vec: np.ndarray | None = None
@@ -220,6 +227,19 @@ def run_cisd(self, *, store_cisd_vec: bool = False) -> None:
         if store_cisd_vec:
             self.cisd_vec = cisd_vec
 
+    def run_sci(self, *, store_sci_vec: bool = False) -> None:
+        """Run SCI and store results."""
+        sci = pyscf.fci.SCI(self.scf.run())
+        sci_energy, sci_vec = sci.kernel(
+            self.one_body_integrals,
+            self.two_body_integrals,
+            norb=self.norb,
+            nelec=self.nelec,
+        )
+        self.sci_energy = sci_energy + self.core_energy
+        if store_sci_vec:
+            self.sci_vec = (sci_vec, *sci_vec._strs)
+
     def run_fci(self, *, store_fci_vec: bool = False) -> None:
         """Run FCI and store results."""
         cas = pyscf.mcscf.CASCI(self.scf.run(), ncas=self.norb, nelecas=self.nelec)
@@ -343,6 +363,8 @@ def as_array_tuple_or_none(val):
             ccsd_t2=arrays_func(data.get("ccsd_t2")),
             cisd_energy=data.get("cisd_energy"),
             cisd_vec=as_array_or_none(data.get("cisd_vec")),
+            sci_energy=data.get("sci_energy"),
+            sci_vec=as_array_tuple_or_none(data.get("sci_vec")),
             fci_energy=data.get("fci_energy"),
             fci_vec=as_array_or_none(data.get("fci_vec")),
             dipole_integrals=as_array_or_none(data.get("dipole_integrals")),
diff --git a/python/ffsim/states/slater.py b/python/ffsim/states/slater.py
index 51f403972..362696e21 100644
--- a/python/ffsim/states/slater.py
+++ b/python/ffsim/states/slater.py
@@ -162,7 +162,6 @@ def _generate_marginals(
     """
     marginals = np.zeros(len(empty_orbitals), dtype=float)
     for i, orbital in enumerate(empty_orbitals):
-        new_sample = sample.copy()
         new_sample = sample + [orbital]
         rest_rdm = rdm[np.ix_(new_sample, new_sample)]
         marginals[i] = np.linalg.det(rest_rdm).real
@@ -189,14 +188,13 @@ def _autoregressive_slater(
         A Numpy array with the position of the sampled electrons.
     """
     rng = np.random.default_rng(seed)
-
     probs = np.diag(rdm).real / nelec
     sample = [rng.choice(norb, p=probs)]
     marginal = [probs[sample[0]]]
     all_orbs = set(range(norb))
     empty_orbitals = list(all_orbs.difference(sample))
     for k in range(nelec - 1):
-        marginals = _generate_marginals(rdm, sample, empty_orbitals)
+        marginals = _generate_marginals(rdm, sample, empty_orbitals, marginal[-1])
         conditionals = marginals / marginal[-1]
         conditionals /= np.sum(conditionals)
         index = rng.choice(len(empty_orbitals), p=conditionals)
@@ -233,12 +231,10 @@ def _sample_spinless_direct(
         Each row is a sample.
     """
     norb, _ = rdm.shape
-
     if nelec == 0:
         return [0] * shots
     if nelec == norb:
         return [(1 << norb) - 1] * shots
     rng = np.random.default_rng(seed)
-    samples = []
     samples = [_autoregressive_slater(rdm, norb, nelec, rng) for _ in range(shots)]
     return [sum(1 << orb for orb in sample) for sample in samples]
diff --git a/python/ffsim/variational/ucj_spin_balanced.py b/python/ffsim/variational/ucj_spin_balanced.py
index 194314f3c..0505c9135 100644
--- a/python/ffsim/variational/ucj_spin_balanced.py
+++ b/python/ffsim/variational/ucj_spin_balanced.py
@@ -410,7 +410,12 @@ def from_t_amplitudes(
         Args:
             t2: The t2 amplitudes.
             t1: The t1 amplitudes.
-            n_reps: The number of ansatz repetitions.
+            n_reps: The number of ansatz repetitions. If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of repetitions.
             interaction_pairs: Optional restrictions on allowed orbital interactions
                 for the diagonal Coulomb operators.
                 If specified, `interaction_pairs` should be a pair of lists,
@@ -450,6 +455,17 @@ def from_t_amplitudes(
         diag_coulomb_mats = np.stack([diag_coulomb_mats, diag_coulomb_mats], axis=1)
         orbital_rotations = orbital_rotations.reshape(-1, norb, norb)[:n_reps]
 
+        n_vecs, _, _, _ = diag_coulomb_mats.shape
+        if n_reps is not None and n_vecs < n_reps:
+            # Pad with no-ops to the requested number of repetitions
+            diag_coulomb_mats = np.concatenate(
+                [diag_coulomb_mats, np.zeros((n_reps - n_vecs, 2, norb, norb))]
+            )
+            eye = np.eye(norb)
+            orbital_rotations = np.concatenate(
+                [orbital_rotations, np.stack([eye for _ in range(n_reps - n_vecs)])]
+            )
+
         final_orbital_rotation = None
         if t1 is not None:
             final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)
diff --git a/python/ffsim/variational/ucj_spin_unbalanced.py b/python/ffsim/variational/ucj_spin_unbalanced.py
index e7e289c74..13aa17ae5 100644
--- a/python/ffsim/variational/ucj_spin_unbalanced.py
+++ b/python/ffsim/variational/ucj_spin_unbalanced.py
@@ -468,6 +468,12 @@ def from_t_amplitudes(
                 the number of terms to use from the alpha-beta t2 amplitudes,
                 and the second integer specifies the number of terms to use from the
                 alpha-alpha and beta-beta t2 amplitudes.
+                If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of repetitions.
             interaction_pairs: Optional restrictions on allowed orbital interactions
                 for the diagonal Coulomb operators.
                 If specified, `interaction_pairs` should be a tuple of 3 lists,
@@ -524,23 +530,27 @@ def from_t_amplitudes(
         diag_coulomb_mats_bb = diag_coulomb_mats_bb.reshape(-1, norb, norb)
         orbital_rotations_bb = orbital_rotations_bb.reshape(-1, norb, norb)
         zero = np.zeros((norb, norb))
-        diag_coulomb_mats_same_spin = np.stack(
-            [
-                [mat_aa, zero, mat_bb]
-                for mat_aa, mat_bb in itertools.zip_longest(
-                    diag_coulomb_mats_aa, diag_coulomb_mats_bb, fillvalue=zero
-                )
-            ]
-        )
-        eye = np.eye(norb)
-        orbital_rotations_same_spin = np.stack(
-            [
-                [orbital_rotation_aa, orbital_rotation_bb]
-                for orbital_rotation_aa, orbital_rotation_bb in itertools.zip_longest(
-                    orbital_rotations_aa, orbital_rotations_bb, fillvalue=eye
-                )
-            ]
-        )
+        if len(diag_coulomb_mats_aa) or len(diag_coulomb_mats_bb):
+            diag_coulomb_mats_same_spin = np.stack(
+                [
+                    [mat_aa, zero, mat_bb]
+                    for mat_aa, mat_bb in itertools.zip_longest(
+                        diag_coulomb_mats_aa, diag_coulomb_mats_bb, fillvalue=zero
+                    )
+                ]
+            )
+            eye = np.eye(norb)
+            orbital_rotations_same_spin = np.stack(
+                [
+                    [orb_rot_aa, orb_rot_bb]
+                    for orb_rot_aa, orb_rot_bb in itertools.zip_longest(
+                        orbital_rotations_aa, orbital_rotations_bb, fillvalue=eye
+                    )
+                ]
+            )
+        else:
+            diag_coulomb_mats_same_spin = np.empty((0, 3, norb, norb))
+            orbital_rotations_same_spin = np.empty((0, 2, norb, norb))
         # concatenate
         if n_reps is None or isinstance(n_reps, int):
             diag_coulomb_mats = np.concatenate(
@@ -549,19 +559,27 @@ def from_t_amplitudes(
             orbital_rotations = np.concatenate(
                 [orbital_rotations_ab, orbital_rotations_same_spin]
             )[:n_reps]
+            diag_coulomb_mats = _pad_diag_coulomb_mats(diag_coulomb_mats, n_reps)
+            orbital_rotations = _pad_orbital_rotations(orbital_rotations, n_reps)
         else:
             n_reps_ab, n_reps_same_spin = n_reps
+            diag_coulomb_mats_ab = _pad_diag_coulomb_mats(
+                diag_coulomb_mats_ab[:n_reps_ab], n_reps_ab
+            )
+            orbital_rotations_ab = _pad_orbital_rotations(
+                orbital_rotations_ab[:n_reps_ab], n_reps_ab
+            )
+            diag_coulomb_mats_same_spin = _pad_diag_coulomb_mats(
+                diag_coulomb_mats_same_spin[:n_reps_same_spin], n_reps_same_spin
+            )
+            orbital_rotations_same_spin = _pad_orbital_rotations(
+                orbital_rotations_same_spin[:n_reps_same_spin], n_reps_same_spin
+            )
             diag_coulomb_mats = np.concatenate(
-                [
-                    diag_coulomb_mats_ab[:n_reps_ab],
-                    diag_coulomb_mats_same_spin[:n_reps_same_spin],
-                ]
+                [diag_coulomb_mats_ab, diag_coulomb_mats_same_spin]
             )
             orbital_rotations = np.concatenate(
-                [
-                    orbital_rotations_ab[:n_reps_ab],
-                    orbital_rotations_same_spin[:n_reps_same_spin],
-                ]
+                [orbital_rotations_ab, orbital_rotations_same_spin]
             )
 
         final_orbital_rotation = None
@@ -676,3 +694,28 @@ def _approx_eq_(self, other, rtol: float, atol: float) -> bool:
                 )
             return True
         return NotImplemented
+
+
+def _pad_diag_coulomb_mats(diag_coulomb_mats: np.ndarray, n_reps: int | None):
+    n_vecs, _, norb, _ = diag_coulomb_mats.shape
+    if n_reps is not None and n_vecs < n_reps:
+        diag_coulomb_mats = np.concatenate(
+            [
+                diag_coulomb_mats,
+                np.zeros((n_reps - n_vecs, 3, norb, norb)),
+            ]
+        )
+    return diag_coulomb_mats
+
+
+def _pad_orbital_rotations(orbital_rotations: np.ndarray, n_reps: int | None):
+    n_vecs, _, norb, _ = orbital_rotations.shape
+    if n_reps is not None and n_vecs < n_reps:
+        eye = np.eye(norb)
+        orbital_rotations = np.concatenate(
+            [
+                orbital_rotations,
+                np.stack([np.stack([eye, eye]) for _ in range(n_reps - n_vecs)]),
+            ]
+        )
+    return orbital_rotations
diff --git a/python/ffsim/variational/ucj_spinless.py b/python/ffsim/variational/ucj_spinless.py
index 40699fc03..4e36fd9b9 100644
--- a/python/ffsim/variational/ucj_spinless.py
+++ b/python/ffsim/variational/ucj_spinless.py
@@ -353,7 +353,12 @@ def from_t_amplitudes(
         Args:
             t2: The t2 amplitudes.
             t1: The t1 amplitudes.
-            n_reps: The number of ansatz repetitions.
+            n_reps: The number of ansatz repetitions. If not specified, then it is set
+                to the number of terms resulting from the double-factorization of the
+                t2 amplitudes. If the specified number of repetitions is larger than the
+                number of terms resulting from the double-factorization, then the ansatz
+                is padded with additional identity operators up to the specified number
+                of repetitions.
             interaction_pairs: Optional restrictions on allowed orbital interactions
                 for the diagonal Coulomb operators.
                 If specified, `interaction_pairs` should be a list of integer pairs
@@ -384,6 +389,17 @@ def from_t_amplitudes(
         diag_coulomb_mats = diag_coulomb_mats.reshape(-1, norb, norb)[:n_reps]
         orbital_rotations = orbital_rotations.reshape(-1, norb, norb)[:n_reps]
 
+        n_vecs, _, _ = diag_coulomb_mats.shape
+        if n_reps is not None and n_vecs < n_reps:
+            # Pad with no-ops to the requested number of repetitions
+            diag_coulomb_mats = np.concatenate(
+                [diag_coulomb_mats, np.zeros((n_reps - n_vecs, norb, norb))]
+            )
+            eye = np.eye(norb)
+            orbital_rotations = np.concatenate(
+                [orbital_rotations, np.stack([eye for _ in range(n_reps - n_vecs)])]
+            )
+
         final_orbital_rotation = None
         if t1 is not None:
             final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)
diff --git a/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py b/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py
index 703bea3de..e483ac3ae 100644
--- a/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py
+++ b/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py
@@ -273,3 +273,24 @@ def test_from_fermion_operator_fermi_hubbard_2d(norb: int, nelec: tuple[int, int
     actual_periodic = actual_linop_periodic @ vec
     expected_periodic = expected_linop_periodic @ vec
     np.testing.assert_allclose(actual_periodic, expected_periodic)
+
+
+def test_diag():
+    """Test computing diagonal."""
+    rng = np.random.default_rng(2222)
+    norb = 5
+    nelec = (3, 2)
+    # TODO test complex one-body after adding support for it
+    one_body_tensor = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mat_a = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mat_b = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mats = np.stack([diag_coulomb_mat_a, diag_coulomb_mat_b])
+    constant = rng.standard_normal()
+    hamiltonian = ffsim.DiagonalCoulombHamiltonian(
+        one_body_tensor, diag_coulomb_mats, constant=constant
+    )
+    linop = ffsim.linear_operator(hamiltonian, norb=norb, nelec=nelec)
+    hamiltonian_dense = linop @ np.eye(ffsim.dim(norb, nelec))
+    np.testing.assert_allclose(
+        ffsim.diag(hamiltonian, norb=norb, nelec=nelec), np.diag(hamiltonian_dense)
+    )
diff --git a/tests/python/molecular_data_test.py b/tests/python/molecular_data_test.py
index 1b9f6bc45..f7aab426e 100644
--- a/tests/python/molecular_data_test.py
+++ b/tests/python/molecular_data_test.py
@@ -31,6 +31,7 @@ def _assert_mol_data_equal(
             "np.ndarray | None",
             "np.ndarray | tuple[np.ndarray, np.ndarray] | None",
             "np.ndarray | tuple[np.ndarray, np.ndarray, np.ndarray] | None",
+            "tuple[np.ndarray, np.ndarray, np.ndarray] | None",
         ]:
             if actual is not None:
                 if isinstance(actual, tuple):
@@ -85,11 +86,13 @@ def test_molecular_data_run_methods():
     mol_data.run_mp2()
     mol_data.run_ccsd()
     mol_data.run_cisd()
+    mol_data.run_sci()
     mol_data.run_fci()
 
     np.testing.assert_allclose(mol_data.mp2_energy, -108.58852784026)
     np.testing.assert_allclose(mol_data.ccsd_energy, -108.5933309085008)
     np.testing.assert_allclose(mol_data.cisd_energy, -108.5878344909782)
+    np.testing.assert_allclose(mol_data.sci_energy, -108.59598682615388)
     np.testing.assert_allclose(mol_data.fci_energy, -108.595987350986)
 
 
@@ -108,6 +111,7 @@ def test_json_closed_shell(tmp_path: pathlib.Path):
     mol_data.run_mp2(store_t2=True)
     mol_data.run_ccsd(store_t1=True, store_t2=True)
     mol_data.run_cisd(store_cisd_vec=True)
+    mol_data.run_sci(store_sci_vec=True)
     mol_data.run_fci(store_fci_vec=True)
 
     for compression in [None, "gzip", "bz2", "lzma"]:
@@ -132,6 +136,7 @@ def test_json_open_shell(tmp_path: pathlib.Path):
     mol_data.run_mp2(store_t2=True)
     mol_data.run_ccsd(store_t1=True, store_t2=True)
     mol_data.run_cisd(store_cisd_vec=True)
+    mol_data.run_sci(store_sci_vec=True)
     mol_data.run_fci(store_fci_vec=True)
 
     for compression in [None, "gzip", "bz2", "lzma"]:
diff --git a/tests/python/states/slater_test.py b/tests/python/states/slater_test.py
index 945b30ca0..00a246847 100644
--- a/tests/python/states/slater_test.py
+++ b/tests/python/states/slater_test.py
@@ -1,4 +1,14 @@
-"""Test Slater sampler."""
+# (C) Copyright IBM 2024.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Tests for Slater determinant functions."""
 
 from __future__ import annotations
 
diff --git a/tests/python/variational/ucj_spin_balanced_test.py b/tests/python/variational/ucj_spin_balanced_test.py
index a4a1615a8..f92efcb71 100644
--- a/tests/python/variational/ucj_spin_balanced_test.py
+++ b/tests/python/variational/ucj_spin_balanced_test.py
@@ -118,10 +118,7 @@ def test_t_amplitudes_energy():
     mol_hamiltonian = mol_data.hamiltonian
 
     # Construct UCJ operator
-    n_reps = 2
-    operator = ffsim.UCJOpSpinBalanced.from_t_amplitudes(
-        ccsd.t2, t1=ccsd.t1, n_reps=n_reps
-    )
+    operator = ffsim.UCJOpSpinBalanced.from_t_amplitudes(ccsd.t2, t1=ccsd.t1)
 
     # Construct the Hartree-Fock state to use as the reference state
     n_alpha, n_beta = nelec
@@ -137,7 +134,45 @@ def test_t_amplitudes_energy():
     # Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state
     hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
     energy = np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state))
-    np.testing.assert_allclose(energy, -108.563917)
+    np.testing.assert_allclose(energy, -108.591373)
+
+
+def test_t_amplitudes_random_n_reps():
+    rng = np.random.default_rng(8379)
+
+    norb = 5
+    nocc = 3
+    nvrt = norb - nocc
+
+    # Construct UCJ operator
+    for n_reps in [3, 15]:
+        t2 = ffsim.random.random_t2_amplitudes(norb, nocc, seed=rng, dtype=float)
+        t1 = rng.standard_normal((nocc, nvrt))
+        operator = ffsim.UCJOpSpinBalanced.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        assert operator.n_reps == n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinBalanced.n_params(
+            norb, n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
+
+
+def test_t_amplitudes_zero_n_reps():
+    norb = 5
+    nocc = 3
+    nvrt = norb - nocc
+
+    # Construct UCJ operator
+    for n_reps in [3, 15]:
+        t2 = np.zeros((nocc, nocc, nvrt, nvrt))
+        t1 = np.zeros((nocc, nvrt))
+        operator = ffsim.UCJOpSpinBalanced.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        assert operator.n_reps == n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinBalanced.n_params(
+            norb, n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
 
 
 def test_t_amplitudes_restrict_indices():
diff --git a/tests/python/variational/ucj_spin_unbalanced_test.py b/tests/python/variational/ucj_spin_unbalanced_test.py
index c6319093b..4e08b530a 100644
--- a/tests/python/variational/ucj_spin_unbalanced_test.py
+++ b/tests/python/variational/ucj_spin_unbalanced_test.py
@@ -109,32 +109,103 @@ def test_t_amplitudes_energy():
     scf = pyscf.scf.ROHF(mol).run()
     ccsd = pyscf.cc.CCSD(scf).run()
 
+    # Get molecular data and molecular Hamiltonian
     mol_data = ffsim.MolecularData.from_scf(scf)
     norb = mol_data.norb
     nelec = mol_data.nelec
     mol_hamiltonian = mol_data.hamiltonian
-    linop = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
 
+    # Construct UCJ operator
+    n_reps = 4
+    operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(
+        ccsd.t2, t1=ccsd.t1, n_reps=n_reps
+    )
+
+    # Construct the Hartree-Fock state to use as the reference state
     n_alpha, n_beta = nelec
     reference_state = ffsim.slater_determinant(
         norb=norb, occupied_orbitals=(range(n_alpha), range(n_beta))
     )
 
+    # Apply the operator to the reference state
+    ansatz_state = ffsim.apply_unitary(
+        reference_state, operator, norb=norb, nelec=nelec
+    )
+
+    # Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state
+    linop = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
+    energy = np.real(np.vdot(ansatz_state, linop @ ansatz_state))
+    np.testing.assert_allclose(energy, -15.124376)
+
+    # Test setting number of reps as tuple
+    n_reps = (4, 2)
     operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(
-        ccsd.t2, t1=ccsd.t1, n_reps=4
+        ccsd.t2, t1=ccsd.t1, n_reps=n_reps
     )
     ansatz_state = ffsim.apply_unitary(
         reference_state, operator, norb=norb, nelec=nelec
     )
     energy = np.real(np.vdot(ansatz_state, linop @ ansatz_state))
-    np.testing.assert_allclose(energy, -15.124376)
+    np.testing.assert_allclose(energy, -15.124832)
 
-    operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(ccsd.t2, n_reps=(4, 2))
+    # Test setting number of reps as None
+    operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(ccsd.t2, t1=ccsd.t1)
     ansatz_state = ffsim.apply_unitary(
         reference_state, operator, norb=norb, nelec=nelec
     )
     energy = np.real(np.vdot(ansatz_state, linop @ ansatz_state))
-    np.testing.assert_allclose(energy, -15.126083)
+    np.testing.assert_allclose(energy, -15.132263)
+
+
+def test_t_amplitudes_random_n_reps():
+    rng = np.random.default_rng(3899)
+    norb = 5
+    nelec = (3, 2)
+    nocc_a, nocc_b = nelec
+    nvrt_a = norb - nocc_a
+    nvrt_b = norb - nocc_b
+
+    for n_reps in [5, 50, (10, 5)]:
+        t2aa = ffsim.random.random_t2_amplitudes(norb, nocc_a, seed=rng, dtype=float)
+        t2ab = rng.standard_normal((nocc_a, nocc_b, nvrt_a, nvrt_b))
+        t2bb = ffsim.random.random_t2_amplitudes(norb, nocc_b, seed=rng, dtype=float)
+        t1a = rng.standard_normal((nocc_a, nvrt_a))
+        t1b = rng.standard_normal((nocc_b, nvrt_b))
+        t2 = (t2aa, t2ab, t2bb)
+        t1 = (t1a, t1b)
+        operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        total_n_reps = n_reps if isinstance(n_reps, int) else sum(n_reps)
+        assert operator.n_reps == total_n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinUnbalanced.n_params(
+            norb, total_n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
+
+
+def test_t_amplitudes_zero_n_reps():
+    norb = 5
+    nelec = (3, 2)
+    nocc_a, nocc_b = nelec
+    nvrt_a = norb - nocc_a
+    nvrt_b = norb - nocc_b
+
+    for n_reps in [5, 50, (10, 5)]:
+        t2aa = np.zeros((nocc_a, nocc_a, nvrt_a, nvrt_a))
+        t2ab = np.zeros((nocc_a, nocc_b, nvrt_a, nvrt_b))
+        t2bb = np.zeros((nocc_b, nocc_b, nvrt_b, nvrt_b))
+        t1a = np.zeros((nocc_a, nvrt_a))
+        t1b = np.zeros((nocc_b, nvrt_b))
+        t2 = (t2aa, t2ab, t2bb)
+        t1 = (t1a, t1b)
+        operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        total_n_reps = n_reps if isinstance(n_reps, int) else sum(n_reps)
+        assert operator.n_reps == total_n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinUnbalanced.n_params(
+            norb, total_n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
 
 
 def test_t_amplitudes_restrict_indices():
diff --git a/tests/python/variational/ucj_spinless_test.py b/tests/python/variational/ucj_spinless_test.py
index 488499046..2ec1ec785 100644
--- a/tests/python/variational/ucj_spinless_test.py
+++ b/tests/python/variational/ucj_spinless_test.py
@@ -113,8 +113,7 @@ def test_t_amplitudes_energy():
     mol_hamiltonian = mol_data.hamiltonian
 
     # Construct UCJ operator
-    n_reps = 2
-    operator = ffsim.UCJOpSpinless.from_t_amplitudes(ccsd.t2, t1=ccsd.t1, n_reps=n_reps)
+    operator = ffsim.UCJOpSpinless.from_t_amplitudes(ccsd.t2, t1=ccsd.t1)
 
     # Compute energy using entanglement forging
     reference_occupations_spatial = [
@@ -143,7 +142,43 @@ def test_t_amplitudes_energy():
         nelec=nelec,
     )
     np.testing.assert_allclose(energy, energy_alt)
-    np.testing.assert_allclose(energy, -108.519714)
+    np.testing.assert_allclose(energy, -108.511945)
+
+
+def test_t_amplitudes_random_n_reps():
+    rng = np.random.default_rng(8379)
+
+    norb = 5
+    nocc = 3
+    nvrt = norb - nocc
+
+    for n_reps in [3, 15]:
+        t2 = ffsim.random.random_t2_amplitudes(norb, nocc, seed=rng, dtype=float)
+        t1 = rng.standard_normal((nocc, nvrt))
+        operator = ffsim.UCJOpSpinless.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        assert operator.n_reps == n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinless.n_params(
+            norb, n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
+
+
+def test_t_amplitudes_zero_n_reps():
+    norb = 5
+    nocc = 3
+    nvrt = norb - nocc
+
+    for n_reps in [3, 15]:
+        t2 = np.zeros((nocc, nocc, nvrt, nvrt))
+        t1 = np.zeros((nocc, nvrt))
+        operator = ffsim.UCJOpSpinless.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)
+        assert operator.n_reps == n_reps
+        actual = len(operator.to_parameters())
+        expected = ffsim.UCJOpSpinless.n_params(
+            norb, n_reps, with_final_orbital_rotation=True
+        )
+        assert actual == expected
 
 
 def test_t_amplitudes_restrict_indices():