code: def threefry_generate(gen, out_shape):
      """Generate a series of random values
  
      Notes
      -----
      This function uses the counter portion of the generator state to generate a series of random
      numbers in parallel. Random number `i` is generated by applying Threefry to the current
      generator state with the counter portion incremented by `i`. This means that each random number
      is generated independently from each other random number, so we can compute them in parallel.
  
      If there is not enough room left in the counter to generate the desired shape of random values,
      then a new generator is created by applying Threefry to the current key, path, and counter.
      This new generator will have a reset counter.
  
      Warning
      -------
      Threeyfry requires that unsigned integer arithmetic wraps on overflow. Currently TVM has no
      guarantee of this, so threefry contains an internal assert to check wrapping behavior. This
      assert may or may not run depending on your platform, so it is recommended you run
      :py:func:`threefry_test_wrapping` to verify wrapping behavior.
  
      Parameters
      ----------
      gen : Tensor[10, uint64]
          Generator state. Can be create with :py:func:`tvm.relay.random.threefry_key`. This should
          not be reused in another function, otherwise random numbers will be repeated.
  
      out_shape : Sequence[int]
          Output shape of the random numbers.
  
      Returns
      -------
      new_gen : Tensor[10, uint64]
          The new generator state to be used in subsequent calls.
  
      rand : Tensor[out_shape, uint64]
          Tensor of random numbers with shape `out_shape`.
      """
      out_len = tir.const(1)
      for s in out_shape:
          out_len *= s
      assert (
          out_len.value <= 2**64 - 1
      ), f"Can only generate up to 2^64 random numbers, but {out_len} were requested."
  
      def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr):
          irb = ir_builder.create()
          gen = irb.buffer_ptr(gen_ptr)
          out_gen = irb.buffer_ptr(out_gen_ptr)
          out_array = irb.buffer_ptr(out_array_ptr)
  
          # Check that unsigned arithmetic wraps, as it is required to implement threefry correctly.
          irb.emit(
              tvm.tir.AssertStmt(
                  tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64")
                  + tvm.tir.const(1, "uint64")
                  == tvm.tir.const(0, "uint64"),
                  tvm.tir.StringImm(
                      "Unsigned integer arithmetic is not wrapping, but threefry requires wrapping."
                  ),
                  tvm.tir.Evaluate(0),
              )
          )
  
          # Create a temporary array to hold the generator state we will use to create the random
          # numbers. We cannot use gen because we may need to update the key + path if there is not
          # enough room in the counter.
          tmp = irb.allocate(gen.dtype, 10, name="tmp", scope="global")
  
          # TODO(tkonolige): for now we only use the last word of the counter for counting. It is too
          # much work to figure out how to do 128 bit addition.
  
          # Max value for counter should be 2**64-2 because we need to reserve a special value to
          # indicate the counter is used up.
          with irb.if_scope(
              gen[7] < tir.const(2**64 - 1, dtype=gen.dtype) - out_len
          ):
              for i in range(10):
                  tmp[i] = gen[i]
          with irb.else_scope():
              # no room left in the counter, we have to change the path or key
              with irb.if_scope(gen[8] == 0 and gen[9] == 0):
                  # out of room in the path, have to generate new key
  
                  # The paper says the counter that we will be hashing should be a special value of
                  # all ones. We need to allocate some space for it because we cannot overwrite gen.
                  tmp_counter = irb.allocate(
                      gen.dtype, 2, name="tmp_counter", scope="global"
                  )
                  tmp_counter[0] = tir.const(
                      0xFFFFFFFFFFFFFFFF, dtype=gen.dtype
                  )
                  tmp_counter[1] = tir.const(
                      0xFFFFFFFFFFFFFFFF, dtype=gen.dtype
                  )
                  _threefry(irb, gen, 0, tmp_counter, 0, tmp, 0, 1)
                  tmp[4] = tir.const(
                      0, dtype=gen.dtype
                  )  # zero path, i.e. no path
                  tmp[5] = tir.const(0, dtype=gen.dtype)
                  tmp[6] = tir.const(0, dtype=gen.dtype)  # zero counter
                  tmp[7] = tir.const(0, dtype=gen.dtype)
                  tmp[8] = tir.const(
                      1 << 63, dtype=gen.dtype
                  )  # one in the leftmost position
                  tmp[9] = tir.const(0, dtype=gen.dtype)
              with irb.else_scope():
                  tmp[0] = gen[0]
                  tmp[1] = gen[1]
                  tmp[2] = gen[2]
                  tmp[3] = gen[3]
                  tmp[4] = gen[4] | gen[8]  # add a 1 to the path
                  tmp[5] = gen[5] | gen[9]
                  tmp[6] = tir.const(0, dtype=gen.dtype)  # zero counter
                  tmp[7] = tir.const(0, dtype=gen.dtype)
                  _shift_right(irb, gen[8], gen[9], tmp, 8, tmp, 9)
  
          # Compute random values
          if out_len.value >= 4:
              _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4)
          if out_len.value % 4 != 0:
              remaining = irb.allocate(
                  gen.dtype, 4, name="remaining", scope="global"
              )
              tmp[7] = tmp[7] + tir.Cast(
                  gen.dtype, out_len // 4 * 4
              )  # increment counter
              _threefry(irb, tmp, 0, tmp, 4, remaining, 0, 1)
              with irb.for_range(
                  0, out_len % 4, dtype="uint64", name="i"
              ) as i:
                  out_array[out_len // 4 * 4 + i] = remaining[i]
  
          # Update generator state
          out_gen[0] = tmp[0]  # key stays the same
          out_gen[1] = tmp[1]
          out_gen[2] = tmp[2]
          out_gen[3] = tmp[3]
          out_gen[4] = tmp[4]  # path stays the same
          out_gen[5] = tmp[5]
          out_gen[6] = tir.const(
              0, dtype=gen.dtype
          )  # unused, leave it as 0
          if out_len.value % 4 != 0:
              # increment counter for the remaining
              # as we will generate 4 random numbers for the remaining, increase 4 here.
              # the main increment was done before the second _threefry.
              out_gen[7] = tmp[7] + tir.Cast(gen.dtype, 4)
          else:
              out_gen[7] = tmp[7] + tir.Cast(
                  gen.dtype, out_len
              )  # increment counter
          out_gen[8] = tmp[8]  # path unchanged, so no update here
          out_gen[9] = tmp[9]
  
          return irb.get()
  
      out_gen = tvm.tir.decl_buffer(
          (10,), name="out_gen", dtype="uint64"
      )
      out_array = tvm.tir.decl_buffer(
          out_shape, name="out_array", dtype="uint64"
      )
      return tvm.te.extern(
          [out_gen.shape, out_array.shape],
          [gen],
          lambda ins, outs: gen_ir(ins[0], outs[0], outs[1]),
          out_buffers=[out_gen, out_array],
          name="threefry_generate",
          tag="threefry_generate",
      )
  
file_path: tvm/python/tvm/topi/random/kernel.py
imports: import math
  import numpy
  import tvm
  import tvm.topi
  from ...None import tir
  from ...tir import ir_builder
label_desc: Use image node_id 2 for calling a global function with example usage: threefry_generate(gen, out_shape) and returns: tvm, irb