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 density matrix function with Tf.
docstr_tok: ["Test","jitting","the","density","matrix","from","density","matrix","function","with","Tf","."]
code: def test_density_matrix_tf_jit(self):
      """Test jitting the density matrix from density matrix function with Tf."""
      from functools import partial
  
      d_mat = tf.Variable(
          [
              [1, 0, 0, 0],
              [0, 0, 0, 0],
              [0, 0, 0, 0],
              [0, 0, 0, 0],
          ],
          dtype=tf.complex128,
      )
      density_matrix = partial(fn.reduce_dm, indices=[0])
  
      density_matrix = tf.function(
          density_matrix,
          jit_compile=True,
          input_signature=(
              tf.TensorSpec(shape=(4, 4), dtype=tf.complex128),
          ),
      )
      density_matrix = density_matrix(d_mat)
      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 10 for calling the TestDensityMatrixFromMatrix obj's underlying member method code with example usage: obj.test_density_matrix_tf_jit() without return types