Coverage for qml_essentials/entanglement.py: 92%

1from typing import Optional, Any, Tuple

2import jax

3import jax.numpy as jnp

4import numpy as np

6import jaqsi as js

7from jaqsi import operations as op

8from jaqsi import gateset

9from jaqsi.math import logm_v

10from qml_essentials.model import Model

11import logging

13log = logging.getLogger(__name__)

16class Entanglement:

17 @classmethod

18 def meyer_wallach(

19 cls,

20 model: Model,

21 n_samples: Optional[int | None],

22 random_key: Optional[jax.random.PRNGKey] = None,

23 scale: bool = False,

24 **kwargs: Any,

25 ) -> float:

26 """

27 Calculates the entangling capacity of a given quantum circuit

28 using Meyer-Wallach measure.

30 Args:

31 model (Model): The quantum circuit model.

32 n_samples (Optional[int]): Number of samples per qubit.

33 If None or < 0, the current parameters of the model are used.

34 random_key (Optional[jax.random.PRNGKey]): JAX random key for

35 parameter initialization. If None, uses the model's internal

36 random key.

37 scale (bool): Whether to scale the number of samples.

38 kwargs (Any): Additional keyword arguments for the model function.

40 Returns:

41 float: Entangling capacity of the given circuit, guaranteed

42 to be between 0.0 and 1.0.

43 """

44 if "noise_params" in kwargs:

45 log.warning(

46 "Meyer-Wallach measure not suitable for noisy circuits. "

47 "Consider 'concentratable entanglement' instead."

48 )

50 if scale:

51 n_samples = jnp.power(2, model.n_qubits) * n_samples

53 if n_samples is not None and n_samples > 0:

54 random_key = model.initialize_params(random_key, repeat=n_samples)

56 # implicitly set input to none in case it's not needed

57 kwargs.setdefault("inputs", None)

58 # explicitly set execution type because everything else won't work

59 rhos = model(execution_type="density", **kwargs).reshape(

60 -1, 2**model.n_qubits, 2**model.n_qubits

61 )

63 ent = cls._compute_meyer_wallach_meas(rhos, model.n_qubits)

65 log.debug(f"Variance of measure: {ent.var()}")

67 return ent.mean()

69 @classmethod

70 def _compute_meyer_wallach_meas(

71 cls, rhos: jnp.ndarray, n_qubits: int

72 ) -> jnp.ndarray:

73 """

74 Computes the Meyer-Wallach entangling capability measure for a given

75 set of density matrices.

77 Args:

78 rhos (jnp.ndarray): Density matrices of the sample quantum states.

79 The shape is (B_s, 2^n, 2^n), where B_s is the number of samples

80 (batch) and n the number of qubits

81 n_qubits (int): The number of qubits

83 Returns:

84 jnp.ndarray: Entangling capability for each sample, array with

85 shape (B_s,)

86 """

87 qb = list(range(n_qubits))

89 def _f(rhos):

90 entropy = 0

91 for j in range(n_qubits):

92 # Formula 6 in https://doi.org/10.48550/arXiv.quant-ph/0305094

93 # Trace out qubit j, keep all others

94 keep = qb[:j] + qb[j + 1 :]

95 density = js.partial_trace(rhos, n_qubits, keep)

96 # only real values, because imaginary part will be separate

97 # in all following calculations anyway

98 # entropy should be 1/2 <= entropy <= 1

99 entropy += jnp.trace((density @ density).real, axis1=-2, axis2=-1)

100

101 # inverse averaged entropy and scale to [0, 1]

102 return 2 * (1 - entropy / n_qubits)

103

104 return jax.vmap(_f)(rhos)

105

106 @classmethod

107 def bell_measurements(

108 cls,

109 model: Model,

110 n_samples: int,

111 random_key: Optional[jax.random.PRNGKey] = None,

112 scale: bool = False,

113 **kwargs: Any,

114 ) -> float:

115 """

116 Compute the Bell measurement for a given model.

117

118 Constructs a ``2 * n_qubits`` circuit that prepares two copies of

119 the model state (on disjoint qubit registers), applies CNOTs and

120 Hadamards, and measures probabilities on the first register.

121

122 Args:

123 model (Model): The quantum circuit model.

124 n_samples (int): The number of samples to compute the measure for.

125 random_key (Optional[jax.random.PRNGKey]): JAX random key for

126 parameter initialization. If None, uses the model's internal

127 random key.

128 scale (bool): Whether to scale the number of samples

129 according to the number of qubits.

130 **kwargs (Any): Additional keyword arguments for the model function.

131

132 Returns:

133 float: The Bell measurement value.

134 """

135 if "noise_params" in kwargs:

136 log.warning(

137 "Bell Measurements not suitable for noisy circuits. "

138 "Consider 'concentratable entanglement' instead."

139 )

140

141 if scale:

142 n_samples = jnp.power(2, model.n_qubits) * n_samples

143

144 n = model.n_qubits

145

146 def _bell_circuit(params, inputs, pulse_params=None, random_key=None, **kw):

147 """Bell measurement circuit on 2*n qubits."""

148 from jaqsi.tape import copy_to_tape

149

150 def vari():

151 model._variational(

152 params,

153 inputs,

154 pulse_params=pulse_params,

155 random_key=random_key,

156 **kw,

157 )

158

159 # First copy on wires 0..n-1

160 vari()

161 # Second copy on wires n..2n-1

162 copy_to_tape(vari, offset=n)

163

164 # Bell measurement: CNOT + H

165 for q in range(n):

166 gateset.CX(wires=[q, q + n])

167 gateset.H(wires=q)

168

169 bell_script = js.Script(f=_bell_circuit, n_qubits=2 * n)

170

171 if n_samples is not None and n_samples > 0:

172 random_key = model.initialize_params(random_key, repeat=n_samples)

173 params = model.params

174 else:

175 if len(model.params.shape) <= 2:

176 params = model.params.reshape(1, *model.params.shape)

177 else:

178 log.info(f"Using sample size of model params: {model.params.shape[0]}")

179 params = model.params

180

181 n_samples = params.shape[0]

182 inputs = model._inputs_validation(kwargs.get("inputs", None))

183

184 # Execute: vmap over batch dimension of params (axis 0)

185 if n_samples > 1:

186 from jaqsi.utils import safe_random_split

187

188 random_keys = safe_random_split(random_key, num=n_samples)

189 result = bell_script.execute(

190 type="probs",

191 args=(params, inputs, model.pulse_params, random_keys),

192 kwargs=kwargs,

193 in_axes=(0, None, None, 0),

194 )

195 else:

196 result = bell_script.execute(

197 type="probs",

198 args=(params, inputs, model.pulse_params, random_key),

199 kwargs=kwargs,

200 )

201

202 # Marginalize: for each qubit q, keep wires [q, q+n] from the 2n-qubit probs

203 # The last probability in each pair gives P(|11⟩) for that qubit pair

204 per_qubit = []

205 for q in range(n):

206 marg = js.marginalize_probs(result, 2 * n, [q, q + n])

207 per_qubit.append(marg)

208 # per_qubit[q] has shape (n_samples, 4) or (4,)

209 exp = jnp.stack(per_qubit, axis=-2) # (..., n, 4)

210 exp = 1 - 2 * exp[..., -1] # (..., n)

211

212 if not jnp.isclose(jnp.sum(exp.imag), 0, atol=1e-6):

213 log.warning("Imaginary part of probabilities detected")

214 exp = jnp.abs(exp)

215

216 measure = 2 * (1 - exp.mean(axis=0))

217 entangling_capability = min(max(float(measure.mean()), 0.0), 1.0)

218 log.debug(f"Variance of measure: {measure.var()}")

219

220 return entangling_capability

221

222 @classmethod

223 def relative_entropy(

224 cls,

225 model: Model,

226 n_samples: int,

227 n_sigmas: int,

228 random_key: Optional[jax.random.PRNGKey] = None,

229 scale: bool = False,

230 **kwargs: Any,

231 ) -> float:

232 """

233 Calculates the relative entropy of entanglement of a given quantum

234 circuit. This measure is also applicable to mixed state, albeit it

235 might me not fully accurate in this simplified case.

236

237 As the relative entropy is generally defined as the smallest relative

238 entropy from the state in question to the set of separable states.

239 However, as computing the nearest separable state is NP-hard, we select

240 n_sigmas of random separable states to compute the distance to, which

241 is not necessarily the nearest. Thus, this measure of entanglement

242 presents an upper limit of entanglement.

243

244 As the relative entropy is not necessarily between zero and one, this

245 function also normalises by the relative entroy to the GHZ state.

246

247 Args:

248 model (Model): The quantum circuit model.

249 n_samples (int): Number of samples per qubit.

250 If <= 0, the current parameters of the model are used.

251 n_sigmas (int): Number of random separable pure states to compare against.

252 random_key (Optional[jax.random.PRNGKey]): JAX random key for

253 parameter initialization. If None, uses the model's internal

254 random key.

255 scale (bool): Whether to scale the number of samples.

256 kwargs (Any): Additional keyword arguments for the model function.

257

258 Returns:

259 float: Entangling capacity of the given circuit, guaranteed

260 to be between 0.0 and 1.0.

261 """

262 dim = jnp.power(2, model.n_qubits)

263 if scale:

264 n_samples = dim * n_samples

265 n_sigmas = dim * n_sigmas

266

267 if random_key is None:

268 random_key = model.random_key

269

270 # Random separable states

271 log_sigmas = sample_random_separable_states(

272 model.n_qubits, n_samples=n_sigmas, random_key=random_key, take_log=True

273 )

274

275 random_key, _ = jax.random.split(random_key)

276

277 if n_samples is not None and n_samples > 0:

278 model.initialize_params(random_key, repeat=n_samples)

279 else:

280 if len(model.params.shape) <= 2:

281 model.params = model.params.reshape(1, *model.params.shape)

282 else:

283 log.info(f"Using sample size of model params: {model.params.shape[0]}")

284

285 rhos, log_rhos = cls._compute_log_density(model, **kwargs)

286

287 rel_entropies = jnp.zeros((n_sigmas, model.params.shape[0]))

288

289 for i, log_sigma in enumerate(log_sigmas):

290 rel_entropies = rel_entropies.at[i].set(

291 cls._compute_rel_entropies(rhos, log_rhos, log_sigma)

292 )

293

294 # Entropy of GHZ states should be maximal

295 ghz_model = Model(model.n_qubits, 1, "GHZ", data_reupload=False)

296 rho_ghz, log_rho_ghz = cls._compute_log_density(ghz_model, **kwargs)

297 ghz_entropies = cls._compute_rel_entropies(rho_ghz, log_rho_ghz, log_sigmas)

298

299 normalised_entropies = rel_entropies / ghz_entropies

300

301 # Average all iterated states

302 entangling_capability = normalised_entropies.T.min(axis=1)

303 log.debug(f"Variance of measure: {entangling_capability.var()}")

304

305 return entangling_capability.mean()

306

307 @classmethod

308 def _compute_log_density(

309 cls, model: Model, **kwargs

310 ) -> Tuple[jnp.ndarray, jnp.ndarray]:

311 """

312 Obtains the density matrix of a model and computes its logarithm.

313

314 Args:

315 model (Model): The model for which to compute the density matrix.

316

317 Returns:

318 Tuple[jnp.ndarray, jnp.ndarray]:

319 - jnp.ndarray: density matrix.

320 - jnp.ndarray: logarithm of the density matrix.

321 """

322 # implicitly set input to none in case it's not needed

323 kwargs.setdefault("inputs", None)

324 # explicitly set execution type because everything else won't work

325 rho = model(execution_type="density", **kwargs)

326 rho = rho.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)

327 log_rho = logm_v(rho) / jnp.log(2)

328 return rho, log_rho

329

330 @classmethod

331 def _compute_rel_entropies(

332 cls,

333 rhos: jnp.ndarray,

334 log_rhos: jnp.ndarray,

335 log_sigmas: jnp.ndarray,

336 ) -> jnp.ndarray:

337 """

338 Compute the relative entropy for a given model.

339

340 Args:

341 rhos (jnp.ndarray): Density matrix result of the circuit, has shape

342 (R, 2^n, 2^n), with the batch size R and number of qubits n

343 log_rhos (jnp.ndarray): Corresponding logarithm of the density

344 matrix, has shape (R, 2^n, 2^n).

345 log_sigmas (jnp.ndarray): Density matrix of next separable state,

346 has shape (2^n, 2^n) if it's a single sigma or (S, 2^n, 2^n),

347 with the batch size S (number of sigmas).

348

349 Returns:

350 jnp.ndarray: Relative Entropy for each sample

351 """

352 n_rhos = rhos.shape[0]

353 if len(log_sigmas.shape) == 3:

354 n_sigmas = log_sigmas.shape[0]

355 rhos = jnp.tile(rhos, (n_sigmas, 1, 1))

356 log_rhos = jnp.tile(log_rhos, (n_sigmas, 1, 1))

357 einsum_subscript = "ij,jk->ik"

358 else:

359 n_sigmas = 1

360 log_sigmas = log_sigmas[jnp.newaxis, ...].repeat(n_rhos, axis=0)

361

362 einsum_subscript = "ij,jk->ik"

363

364 def _f(rhos, log_rhos, log_sigmas):

365 prod = jnp.einsum(einsum_subscript, rhos, log_rhos - log_sigmas)

366 rel_entropies = jnp.abs(jnp.trace(prod, axis1=-2, axis2=-1))

367 return rel_entropies

368

369 rel_entropies = jax.vmap(_f, in_axes=(0, 0, 0))(rhos, log_rhos, log_sigmas)

370

371 if n_sigmas > 1:

372 rel_entropies = rel_entropies.reshape(n_sigmas, n_rhos)

373 return rel_entropies

374

375 @classmethod

376 def entanglement_of_formation(

377 cls,

378 model: Model,

379 n_samples: int,

380 random_key: Optional[jax.random.PRNGKey] = None,

381 scale: bool = False,

382 always_decompose: bool = False,

383 **kwargs: Any,

384 ) -> float:

385 """

386 This function implements the entanglement of formation for mixed

387 quantum systems.

388 In that a mixed state gets decomposed into pure states with respective

389 probabilities using the eigendecomposition of the density matrix.

390 Then, the Meyer-Wallach measure is computed for each pure state,

391 weighted by the eigenvalue.

392 See e.g. https://doi.org/10.48550/arXiv.quant-ph/0504163

393

394 Note that the decomposition is *not unique*! Therefore, this measure

395 presents the entanglement for *some* decomposition into pure states,

396 not necessarily the one that is anticipated when applying the Kraus

397 channels.

398 If a pure state is provided, this results in the same value as the

399 Entanglement.meyer_wallach function if `always_decompose` flag is not set.

400

401 Args:

402 model (Model): The quantum circuit model.

403 n_samples (int): Number of samples per qubit.

404 random_key (Optional[jax.random.PRNGKey]): JAX random key for

405 parameter initialization. If None, uses the model's internal

406 random key.

407 scale (bool): Whether to scale the number of samples.

408 always_decompose (bool): Whether to explicitly compute the

409 entantlement of formation for the eigendecomposition of a pure

410 state.

411 kwargs (Any): Additional keyword arguments for the model function.

412

413 Returns:

414 float: Entangling capacity of the given circuit, guaranteed

415 to be between 0.0 and 1.0.

416 """

417

418 if scale:

419 n_samples = jnp.power(2, model.n_qubits) * n_samples

420

421 if n_samples is not None and n_samples > 0:

422 model.initialize_params(random_key, repeat=n_samples)

423 else:

424 if len(model.params.shape) <= 2:

425 model.params = model.params.reshape(1, *model.params.shape)

426 else:

427 log.info(f"Using sample size of model params: {model.params.shape[0]}")

428

429 # implicitly set input to none in case it's not needed

430 kwargs.setdefault("inputs", None)

431 rhos = model(execution_type="density", **kwargs)

432 rhos = rhos.reshape(-1, 2**model.n_qubits, 2**model.n_qubits)

433 ent = cls._compute_entanglement_of_formation(

434 rhos, model.n_qubits, always_decompose

435 )

436 return ent.mean()

437

438 @classmethod

439 def _compute_entanglement_of_formation(

440 cls,

441 rhos: jnp.ndarray,

442 n_qubits: int,

443 always_decompose: bool,

444 ) -> jnp.ndarray:

445 """

446 Computes the entanglement of formation for a given batch of density

447 matrices.

448

449 Args:

450 rho (jnp.ndarray): The density matrices, has shape (B_s, 2^n, 2^n),

451 where B_s is the batch size and n the number of qubits.

452 n_qubits (int): Number of qubits

453 always_decompose (bool): Whether to explicitly compute the

454 entantlement of formation for the eigendecomposition of a pure

455 state.

456

457 Returns:

458 jnp.ndarray: Entanglement for the provided density matrices.

459 """

460 eigenvalues, eigenvectors = jnp.linalg.eigh(rhos)

461 if not always_decompose and jnp.isclose(eigenvalues, 1.0).any(axis=-1).all():

462 return cls._compute_meyer_wallach_meas(rhos, n_qubits)

463

464 rhos = np.einsum("sij,sik->sijk", eigenvectors, eigenvectors.conjugate())

465 measures = cls._compute_meyer_wallach_meas(

466 rhos.reshape(-1, 2**n_qubits, 2**n_qubits), n_qubits

467 )

468 ent = np.einsum("si,si->s", measures.reshape(-1, 2**n_qubits), eigenvalues)

469 return ent

470

471 @classmethod

472 def concentratable_entanglement(

473 cls,

474 model: Model,

475 n_samples: int,

476 random_key: Optional[jax.random.PRNGKey] = None,

477 scale: bool = False,

478 **kwargs: Any,

479 ) -> float:

480 """

481 Computes the concentratable entanglement of a given model.

482

483 This method utilizes the Concentratable Entanglement measure from

484 https://arxiv.org/abs/2104.06923. The swap test is implemented

485 directly in jaqsi using a ``3 * n_qubits`` circuit.

486

487 Args:

488 model (Model): The quantum circuit model.

489 n_samples (int): The number of samples to compute the measure for.

490 random_key (Optional[jax.random.PRNGKey]): JAX random key for

491 parameter initialization. If None, uses the model's internal

492 random key.

493 scale (bool): Whether to scale the number of samples according to

494 the number of qubits.

495 **kwargs (Any): Additional keyword arguments for the model function.

496

497 Returns:

498 float: Entangling capability of the given circuit, guaranteed

499 to be between 0.0 and 1.0.

500 """

501 n = model.n_qubits

502 N = 2**n

503

504 if scale:

505 n_samples = N * n_samples

506

507 def _swap_test_circuit(

508 params, inputs, pulse_params=None, random_key=None, **kw

509 ):

510 """Swap-test circuit on 3*n qubits."""

511 from jaqsi.tape import copy_to_tape

512

513 def vari():

514 model._variational(

515 params,

516 inputs,

517 pulse_params=pulse_params,

518 random_key=random_key,

519 **kw,

520 )

521

522 # First copy on wires n..2n-1

523 copy_to_tape(vari, offset=n)

524 # Second copy on wires 2n..3n-1

525 copy_to_tape(vari, offset=2 * n)

526

527 # Swap test: H on ancilla register (wires 0..n-1)

528 for i in range(n):

529 gateset.H(wires=i)

530

531 for i in range(n):

532 gateset.CSWAP(wires=[i, i + n, i + 2 * n])

533

534 for i in range(n):

535 gateset.H(wires=i)

536

537 swap_script = js.Script(f=_swap_test_circuit, n_qubits=3 * n)

538

539 if n_samples is not None and n_samples > 0:

540 random_key = model.initialize_params(random_key, repeat=n_samples)

541 else:

542 if len(model.params.shape) <= 2:

543 model.params = model.params.reshape(1, *model.params.shape)

544 else:

545 log.info(f"Using sample size of model params: {model.params.shape[0]}")

546

547 params = model.params

548 inputs = model._inputs_validation(kwargs.get("inputs", None))

549 n_batch = params.shape[0]

550

551 marg_probs = jax.jit(js.marginalize_probs, static_argnums=(1, 2))

552

553 if n_batch > 1:

554 from jaqsi.utils import safe_random_split

555

556 random_keys = safe_random_split(random_key, num=n_batch)

557 probs = swap_script.execute(

558 type="probs",

559 args=(params, inputs, model.pulse_params, random_keys),

560 in_axes=(0, None, None, 0),

561 kwargs=kwargs,

562 )

563 else:

564 probs = swap_script.execute(

565 type="probs",

566 args=(params, inputs, model.pulse_params, random_key),

567 kwargs=kwargs,

568 )

569

570 # Marginalize to the ancilla register (wires 0..n-1)

571 probs = marg_probs(probs, 3 * n, tuple(range(n)))

572

573 ent = 1 - probs[..., 0]

574

575 log.debug(f"Variance of measure: {ent.var()}")

576

577 return float(ent.mean())

578

579 @classmethod

580 def concentratable_entanglement_estimation(

581 cls,

582 model: Model,

583 n_samples: int,

584 random_key: Optional[jax.random.PRNGKey] = None,

585 scale: bool = False,

586 **kwargs: Any,

587 ) -> float:

588 """

589 Computes the concentratable entanglement of a given model.

590

591 This method utilizes the Concentratable Entanglement measure from

592 https://arxiv.org/abs/2104.06923. The swap test is implemented

593 directly in jaqsi using a ``3 * n_qubits`` circuit.

594

595 Args:

596 model (Model): The quantum circuit model.

597 n_samples (int): The number of samples to compute the measure for.

598 random_key (Optional[jax.random.PRNGKey]): JAX random key for

599 parameter initialization. If None, uses the model's internal

600 random key.

601 scale (bool): Whether to scale the number of samples according to

602 the number of qubits.

603 **kwargs (Any): Additional keyword arguments for the model function.

604

605 Returns:

606 float: Entangling capability of the given circuit, guaranteed

607 to be between 0.0 and 1.0.

608 """

609 n = model.n_qubits

610 N = 2**n

611

612 if scale:

613 n_samples = N * n_samples

614

615 def _bell_basis_measurement(

616 params, inputs, pulse_params=None, random_key=None, **kw

617 ):

618 """Bell-basis measurement circuit on 3*n qubits."""

619 from jaqsi.tape import copy_to_tape

620

621 def vari():

622 model._variational(

623 params,

624 inputs,

625 pulse_params=pulse_params,

626 random_key=random_key,

627 **kw,

628 )

629

630 # First copy on wires 0..n-1

631 copy_to_tape(vari, offset=0)

632 # Second copy on wires n..2n-1

633 copy_to_tape(vari, offset=n)

634

635 for i in range(n):

636 gateset.CX(wires=[i, i + n])

637 gateset.H(wires=i)

638

639 bell_basis_script = js.Script(f=_bell_basis_measurement, n_qubits=2 * n)

640

641 if n_samples is not None and n_samples > 0:

642 random_key = model.initialize_params(random_key, repeat=n_samples)

643 else:

644 if len(model.params.shape) <= 2:

645 model.params = model.params.reshape(1, *model.params.shape)

646 else:

647 log.info(f"Using sample size of model params: {model.params.shape[0]}")

648

649 params = model.params

650 inputs = model._inputs_validation(kwargs.get("inputs", None))

651 n_batch = params.shape[0]

652

653 # SWAP operator in Bell-basis

654 SWAP = jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, -1]])

655 # Construct observable for measuring CE

656 CE_observable = gateset.Id([0, n]) + op.Operation([0, n], SWAP)

657 for i in range(1, n):

658 CE_observable = CE_observable @ (

659 gateset.Id([i, i + n]) + op.Operation([i, i + n], SWAP)

660 )

661 CE_observable = (1 / N) * CE_observable

662

663 expvals = []

664 if n_batch > 1:

665 from jaqsi.utils import safe_random_split

666

667 random_keys = safe_random_split(random_key, num=n_batch)

668 expvals = bell_basis_script.execute(

669 type="expval",

670 obs=[CE_observable],

671 args=(params, inputs, model.pulse_params, random_keys),

672 in_axes=(0, None, None, 0),

673 kwargs=kwargs,

674 )

675 else:

676 expvals = bell_basis_script.execute(

677 type="expval",

678 obs=[CE_observable],

679 args=(params, inputs, model.pulse_params, random_key),

680 kwargs=kwargs,

681 )

682

683 ent = 1 - expvals

684 log.debug(f"Variance of measure: {ent.var()}")

685 return float(ent.mean())

686

687

688def sample_random_separable_states(

689 n_qubits: int,

690 n_samples: int,

691 random_key: jax.random.PRNGKey,

692 take_log: bool = False,

693) -> jnp.ndarray:

694 """

695 Sample random separable states (density matrix).

696

697 Args:

698 n_qubits (int): number of qubits in the state

699 n_samples (int): number of states

700 random_key (random.PRNGKey): JAX random key

701 take_log (bool): if the matrix logarithm of the density matrix should be taken.

702

703 Returns:

704 jnp.ndarray: Density matrices of shape (n_samples, 2**n_qubits, 2**n_qubits)

705 """

706 model = Model(n_qubits, 1, "No_Entangling", data_reupload=False)

707 model.initialize_params(random_key, repeat=n_samples)

708 # explicitly set execution type because anything else won't work

709 sigmas = model(execution_type="density", inputs=None)

710 if take_log:

711 sigmas = logm_v(sigmas) / jnp.log(2.0 + 0j)

712

713 return sigmas

Coverage for qml_essentials / entanglement.py: 92%

237 statements