code: def _identity_matrix(
      num_columns: types.IntTensor,
      num_digits: types.IntTensor,
      dtype: tf.DType = None,
  ) -> types.IntTensor:
      r"""Returns the identity matrix.
  
      Args:
        num_columns: Positive scalar `Tensor` with rank 0 representing the number of
          columns of the returned matrix.
        num_digits: Positive scalar `Tensor` with rank 0 representing the base-2
          precision of the samples.
        dtype: Optional `dtype`. The `dtype` of the output `Tensor` (either a signed
          or unsigned integer `dtype`).
          Default value: `None` which maps to `int32`.
  
      Returns:
        A scalar `Tensor` with shape `(1, num_columns)`.
      """
  
      dtype = dtype or tf.int32
      shifts = tf.range(
          num_digits - 1, num_digits - 1 - num_columns, delta=-1
      )
      # shape: (1, num_columns)
      return tf.bitwise.left_shift(
          tf.ones(shape=(1, num_columns), dtype=dtype),
          tf.cast(shifts, dtype),
      )
  
file_path: tf-quant-finance/tf_quant_finance/math/qmc/sobol.py
imports: import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance.math.qmc import digital_net
  from tf_quant_finance.math.qmc import utils
  from tf_quant_finance.math.random_ops import sobol
label_desc: Use image node_id 3 for calling a global function with example usage: _identity_matrix(num_columns, num_digits, dtype) and returns: tf