How to use the trax.tf_numpy.numpy.utils.tensor_to_ndarray function in trax

# Algorithm:
  # 1. Calculate cumulative sum of repeats.
  repeats_cumsum = cumsum(repeats)  # [3, 4, 6]
  # 2. Use `scatter_nd` to generate an indices list for use in `tf.gather`.
  scatter_indices_t = repeats_cumsum[:-1].data  # [3, 4]
  scatter_indices_t = tf.expand_dims(scatter_indices_t, 1)  # [[3], [4]]
  scatter_updates_t = tf.ones([len(repeats) - 1], dtype=tf.int32)  # [1, 1]
  scatter_shape_t = ravel(repeats_cumsum[-1]).data  # [6]
  #    `tf.scatter_nd([[3], [4]], [1, 1], [6])` -> `[0, 0, 0, 1, 1, 0]`
  indices_t = tf.scatter_nd(scatter_indices_t, scatter_updates_t,
                            scatter_shape_t)
  indices_t = tf.cumsum(indices_t)  # [0, 0, 0, 1, 2, 2]
  # 3. Use `tf.gather` to gather indices along `axis`.
  result_t = tf.gather(a, indices_t, axis=axis)

  return utils.tensor_to_ndarray(result_t)

to power of `base`.
    num: Number of values to sample. Defaults to 50.
    endpoint: When to include `base ** stop` in the output. Defaults to true.
    base: Base of the log space.
    dtype: Optional. Type of the resulting ndarray. Could be a python type, a
      NumPy type or a TensorFlow `DType`. If not provided, it is figured from
      input args.
  """
  # TODO(srbs): Check whether dtype is handled properly.
  if dtype:
    dtype = utils.to_tf_type(dtype)
  result = linspace(start, stop, num=num, endpoint=endpoint)
  result = tf.pow(base, result.data)
  if dtype:
    result = utils.maybe_cast(result, dtype)
  return utils.tensor_to_ndarray(result)

This relies on `tf.tensordot` which does not support types int64 and float64.
  So arrays of those types are "unsafely" cast to int32 and float32.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    b: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.

  Returns:
    An ndarray.
  """
  a, b = promote_args_types(a, b)
  if utils.isscalar(a) or utils.isscalar(b):
    a = utils.tensor_to_ndarray(tf.expand_dims(a.data, -1))
    b = utils.tensor_to_ndarray(tf.expand_dims(b.data, -1))
    a_axis = b_axis = -1
  else:
    a_axis = -1
    # TODO(agarwal): handle ndim being None when in graph mode.
    b_axis = -2 if b.ndim > 1 else -1
  # TODO(srbs): When the shape of the output is a scalar e.g. when performing
  # a dot-product of two vectors, numpy returns a scalar object and not an
  # instance of ndarray.

  # tensordot/MatMul does not support int64 and float64 so we manually cast to
  # the compatible types. The conversion may be unsafe.
  # TODO(srbs): Figure out why MatMul does not support larger types.
  output_type = None
  if a.dtype == np.int64:
    logging.warning('Unsafe cast to int32.')

"""Permutes dimensions of the array.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    axes: array_like. A list of ints with length rank(a) or None specifying the
      order of permutation. The i'th dimension of the output array corresponds
      to axes[i]'th dimension of the `a`. If None, the axes are reversed.

  Returns:
    An ndarray.
  """
  a = array_creation.asarray(a)
  if axes is not None:
    axes = array_creation.asarray(axes)
  return utils.tensor_to_ndarray(tf.transpose(a=a.data, perm=axes))

raise ValueError('stop: {} must be non-zero.'.format(stop))
  if np_sign(start) != np_sign(stop):
    raise ValueError('start: {} and stop: {} must have same sign.'.format(
        start, stop))
  step = 1.
  if endpoint:
    if num > 1:
      step = tf.pow((stop / start), 1 / (num - 1))
  else:
    step = tf.pow((stop / start), 1 / num)
  result = tf.cast(tf.range(num), step.dtype)
  result = tf.pow(step, result)
  result = tf.multiply(result, start)
  if dtype:
    result = tf.cast(result, dtype=dtype)
  return utils.tensor_to_ndarray(result)

def squeeze(a, axis=None):
  """Removes single-element axes from the array.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    axis: scalar or list/tuple of ints.

  TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution
  is enabled. So we cannot allow axis to be array_like here. Fix.

  Returns:
    An ndarray.
  """
  a = array_creation.asarray(a)
  return utils.tensor_to_ndarray(tf.squeeze(a, axis))

to the other.

  Args:
    x1: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    x2: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.

  Returns:
    An ndarray of type bool and broadcasted shape of x1 and x2.
  """
  dtype = utils.result_type(x1, x2)
  # Cast x1 and x2 to the result_type if needed.
  x1 = array_creation.array(x1, copy=False, dtype=dtype)
  x2 = array_creation.array(x2, copy=False, dtype=dtype)
  return utils.tensor_to_ndarray(tf.equal(x1.data, x2.data))

If `a` is already a 1-d ndarray it is returned as is.

  Uses `tf.reshape`.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.

  Returns:
    A 1-d ndarray.
  """
  a = array_creation.asarray(a)
  if a.ndim == 1:
    return a
  return utils.tensor_to_ndarray(tf.reshape(a.data, [-1]))

"""Computes the tf_fn(x) for each element in `x`.

  Args:
    x: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    tf_fn: function that takes a single Tensor argument.

  Returns:
    An ndarray with the same shape as `x`. The default output dtype is
    determined by `dtypes.default_float_type`, unless x is an ndarray with a
    floating point type, in which case the output type is same as x.dtype.
  """
  x = array_creation.asarray(x)
  if x.dtype not in (np.float16, np.float32, np.float64):
    x = x.astype(dtypes.default_float_type())
  return utils.tensor_to_ndarray(tf_fn(x.data))

"""Computes the max of all array elements or along specified axes.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    axis: Optional 0-d or 1-d array_like. Axes along which to compute the max.
      If None, returns the max of all elements in array.
    keepdims: If true, retains reduced dimensions with length 1.

  Returns:
    An ndarray with the same dtype as `a`.
  """
  # TODO(wangpeng): check that we fully match numpy behavior.
  a = array_creation.asarray(a)

  return utils.tensor_to_ndarray(tf.reduce_max(input_tensor=a.data, axis=axis,
                                               keepdims=keepdims))