code: def ids_tensor(shape, vocab_size, rng=None):
      """Creates a random int32 tensor of the shape within the vocab size."""
      if rng is None:
          rng = random.Random()
  
      total_dims = 1
      for dim in shape:
          total_dims *= dim
  
      values = []
      for _ in range(total_dims):
          values.append(rng.randint(0, vocab_size - 1))
  
      output = np.array(values, dtype=jnp.int32).reshape(shape)
  
      return output
  
imports: import random
  import unittest
  import numpy
  import transformers
  from transformers import is_flax_available, is_torch_available
  from transformers.testing_utils import is_pt_flax_cross_test, require_flax
label_desc: Use image node_id 1 for calling a global function with example usage: ids_tensor(shape, vocab_size, rng) and returns: output