class_docstr: Tests for the (reduced) density matrix for matrix.
class_docstr_tok: ["Tests","for","the","(","reduced",")","density","matrix","for","matrix","."]
docstr: Test jitting the density matrix from state vector function.
docstr_tok: ["Test","jitting","the","density","matrix","from","state","vector","function","."]
code: def test_reduce_statevector_jax_jit(self):
      """Test jitting the density matrix from state vector function."""
      from jax import jit
  
      state_vector = jnp.array(
          [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
      )
  
      jitted_dens_matrix_func = jit(
          fn.quantum.reduce_dm, static_argnums=[1, 2]
      )
  
      density_matrix = jitted_dens_matrix_func(
          state_vector, indices=(0,), check_state=True
      )
      assert np.allclose(density_matrix, [[1, 0], [0, 0]])
  
imports: import numpy
  import pytest
  from pennylane import numpy
  from pennylane import math
label_desc: Use image node_id 8 for calling the TestDensityMatrixFromMatrix obj's underlying member method code with example usage: obj.test_reduce_statevector_jax_jit() without return types