code: def _sobol_generating_matrices(
      dim: types.IntTensor,
      log_num_results: types.IntTensor,
      num_digits: types.IntTensor,
      dtype=None,
  ) -> types.IntTensor:
      r"""Returns all Sobol generating matrices.
  
      Args:
        dim: Positive scalar `Tensor` with rank 0 representing the event size of
          points which can be sampled from the resulting generating matrix.
        log_num_results: Positive scalar `Tensor` with rank 0 representing the
          base-2 logarithm of the maximum number of points which can be sampled from
          the resulting generating matrix.
        num_digits: Positive scalar `Tensor` with rank 0 representing the base-2
          precision of points which can be sampled from the resulting generating
          matrix.
        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 `(dim, ceil(log2(num_results)))`.
      """
      global _INITIAL_DIRECTION_NUMBERS
      global _PRIMITIVE_POLYNOMIAL_COEFFICIENTS
  
      dtype = dtype or tf.int32
  
      zero = tf.constant(0, dtype=dtype)
      indices = tf.cast(tf.range(0, log_num_results), dtype)
      dimensions = tf.range(0, dim)
  
      # shape: (?, ?)
      directions = tf.convert_to_tensor(
          _INITIAL_DIRECTION_NUMBERS,
          dtype=dtype,
          name="direction_numbers",
      )
      padding = log_num_results - utils.get_shape(directions)[0]
      padding = tf.math.maximum(zero, padding)
      directions = tf.pad(directions, [[zero, padding], [zero, zero]])
      # shape: (log_num_results, ?)
      directions = directions[:log_num_results]
      # shape: (log_num_results, dim)
      directions = tf.gather(directions, dimensions, axis=1)
      # shape: (dim, log_num_results)
      directions = tf.cast(tf.transpose(directions), dtype)
  
      # shape: (?,)
      polynomial = tf.convert_to_tensor(
          _PRIMITIVE_POLYNOMIAL_COEFFICIENTS,
          dtype=dtype,
          name="polynomial_coefficients",
      )
      # shape: (1, dim)
      polynomial = tf.cast(
          tf.gather(polynomial, tf.expand_dims(dimensions, axis=1)),
          dtype,
      )
  
      # shape: (1, dim)
      degree = tf.cast(
          tf.math.floor(
              utils.log2(tf.cast(polynomial, dtype=tf.float32))
          ),
          dtype=dtype,
      )
  
      # The values of the `i`-th generating matrix (i.e. the `i`-th row of the
      # output tensor) are calculated using the formulas below.
  
      # If `col < degree`, the value is obtained from the direction numbers:
      #
      #   matrix[col] = directions[col] << num_digits - 1 - col
      #
      # where:
      #   `directions` are the direction numbers.
  
      # shape: (dim, log_num_results)
      initial_matrices = tf.bitwise.left_shift(
          directions,
          tf.cast(
              tf.expand_dims(num_digits - 1 - indices, axis=0), dtype
          ),
      )
  
      # If `degree <= col < log_num_results`, the value is obtained using the
      # following recurrence:
      #
      #   matrix[col] = (matrix[col - degree] >> degree)
      #               ^ (matrix[col - (degree - 0)] * (polynomial >> 0) & 1))
      #               ^ (matrix[col - (degree - 1]) * (polynomial >> 1) & 1))
      #               ^ (matrix[col - (degree - 2]) * (polynomial >> 2) & 1))
      #               ^ ...
      #               ^ (matrix[col - 2] * (polynomial >> (degree - 2)) & 1))
      #               ^ (matrix[col - 1] * (polynomial >> (degree - 1)) & 1))
      #
      # where:
      #  `polynomial` is the `i`-th primitive polynomial, and `degree` its degree.
      #
      # This recurrence normally requires three nested loops:
      # - One on the matrix index (i.e. `i`);
      # - One on the matrix column (i.e. `col`);
      # - One for the recurrence.
      #
      # The number of iterations applied by the innermost loop depends on both `i`
      # (through `degree`) and `col`. These dependencies must be removed in order
      # to vectorize the calculation on GPUs.
      #
      # For this, we relax the constraint on `col` and apply the recurrence for
      # `col < log_num_results`. We then add this constraint back by filtering
      # out values which should not be updated with `tf.where`. With this approach,
      # the algorithm can be reduced to a single explicit while-loop on `col`. Some
      # calculations are wasted since their results are thrown away, but those are
      # essentially free on GPUs due to vectorization.
      #
      # Note that `column` in the implementation below does not correspond to `col`.
      # Instead, it is the innermost loop index and ranges from `col - degree` to
      # `col - 1` after filtering.
  
      def loop_predicate_fn(matrix_values, column):
          del matrix_values
          return column < log_num_results - 1
  
      # Loop invariant:
      #   At the end of the iteration for column `column`, values for columns `0` to
      #   `column` have been calculated. Values for other columns are only partially
      #   calculated.
      def loop_body_fn(matrices, column):
          # shape: (dim, log_num_results)
          column_values = tf.gather(matrices, [column], axis=1)
  
          # shape: (dim, log_num_results)
          should_be_updated = tf.logical_and(
              # Columns whose index is smaller than the degree of the primitive
              # polynomial are obtained from direction numbers and should not be
              # updated.
              tf.less_equal(
                  tf.math.maximum(degree, column + 1), indices
              ),
              # During a given iteration, only the next `n` columns (where `n` is the
              # degree of the primitive polynomial) should be updated.
              tf.less_equal(indices, column + degree),
          )
  
          # shape: (dim, log_num_results)
          updated_matrices = tf.bitwise.bitwise_xor(
              tf.where(
                  tf.equal(indices, column + degree),
                  tf.bitwise.right_shift(column_values, degree),
                  matrices,
              ),
              utils.filter_tensor(
                  column_values, polynomial, column + degree - indices
              ),
          )
  
          # shape: (dim, log_num_results)
          returned_matrices = tf.where(
              should_be_updated, updated_matrices, matrices
          )
  
          return (returned_matrices, column + 1)
  
      matrices, _ = tf.while_loop(
          loop_predicate_fn,
          loop_body_fn,
          loop_vars=(initial_matrices, tf.constant(0, dtype)),
          maximum_iterations=tf.cast(log_num_results, tf.int32) - 1,
      )
  
      # shape: (dim, log_num_results)
      return matrices
  
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 4 for calling a global function with example usage: _sobol_generating_matrices(dim, log_num_results, num_digits, dtype) and returns: matrices, unknown, returned_matrices,  unknown