Python tensorflow.compat.v2.TensorShape() Examples

def _reshape_inner_dims(
    tensor: tf.Tensor,
    shape: tf.TensorShape,
    new_shape: tf.TensorShape) -> tf.Tensor:
  """Reshapes tensor to: shape(tensor)[:-len(shape)] + new_shape."""
  tensor_shape = tf.shape(tensor)
  ndims = shape.rank
  tensor.shape[-ndims:].assert_is_compatible_with(shape)
  new_shape_inner_tensor = tf.cast(
      [-1 if d is None else d for d in new_shape.as_list()], tf.int64)
  new_shape_outer_tensor = tf.cast(
      tensor_shape[:-ndims], tf.int64)
  full_new_shape = tf.concat(
      (new_shape_outer_tensor, new_shape_inner_tensor), axis=0)
  new_tensor = tf.reshape(tensor, full_new_shape)
  new_tensor.set_shape(tensor.shape[:-ndims] + new_shape)
  return new_tensor 

def spec_as_shape(spec, context):
  """Convert a type_spec to a tf shape.

  Args:
    spec: a single specification for tuple_generator_builder
    context: a NQL context

  Returns:
    tensor shape specification, as required by tf.data.Dataset.from_generator
  """
  if spec == str:
    return tf.TensorShape([])
  elif isinstance(spec, int):
    return tf.TensorShape([spec])
  else:
    return tf.TensorShape([context.get_max_id(spec)])


# GOOGLE_INTERNAL: TODO(b/124102056) Consider moving into nql. 

def for_with_variable_shape_growing_matrix_rows(l):
  m = tf.constant([[0]])
  for i in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(m, tf.TensorShape([None, 1]))])
    m = tf.concat((m, [[i]]), 0)
  return m 

def _event_shape(self):
    return tf.TensorShape(()) 

def _batch_shape(self):
    return tf.TensorShape(self._batch_shape_tuple) 

def _check_value(self, tensor: tf.Tensor, tensorspec: tf.TensorSpec):
    if not tf.TensorShape(tf.squeeze(tensor.get_shape())).is_compatible_with(
        tensorspec.shape):
      raise ValueError(
          'Tensor {} is not compatible with specification {}.'.format(
              tensor, tensorspec)) 

def for_with_shape_invariant_violation(l):
  t = tf.constant([1])
  for _ in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=((t, tf.TensorShape([1])),))
    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
  return t 

def while_with_shape_invariant_violation():
  t = tf.constant([1])
  while tf.constant(True):
    tf.autograph.experimental.set_loop_options(
        shape_invariants=((t, tf.TensorShape([1])),))
    t = tf.range(tf.random.uniform((), 2, 3, dtype=tf.int32))
  return t 

def for_with_composite_tensor_shape_invariant(l):
  v = tf.SparseTensor(
      indices=[[0, 0], [1, 1]], values=[1, 2], dense_shape=[3, 3])
  for _ in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape(None))])
    v = tf.sparse.expand_dims(v)
  return v 

def while_with_composite_tensor_shape_invariant(n):
  v = tf.SparseTensor(
      indices=[[0, 0], [1, 1]], values=[1, 2], dense_shape=[3, 3])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape(None))])
    v = tf.sparse.expand_dims(v)
    i += 1
  return v 

def while_with_variable_shape_and_break(n):
  v = tf.constant([0, 0])
  i = 0
  if n > 1:
    while i < n:
      tf.autograph.experimental.set_loop_options(
          shape_invariants=[(v, tf.TensorShape([None]))])
      v = tf.concat((v, [i]), 0)
      i += 1
      if i > 3:
        break
  else:
    v = tf.constant([1, 2, 3])
  return v 

def for_with_variable_shape_inside_if(n):
  v = tf.constant([0, 0])
  if n > 1:
    for i in range(n):
      tf.autograph.experimental.set_loop_options(
          shape_invariants=[(v, tf.TensorShape([None]))])
      v = tf.concat((v, [i]), 0)
      i += 1
  else:
    v = tf.constant([1, 2, 3])
  return v 

def while_with_variable_shape_inside_if(n):
  v = tf.constant([0, 0])
  i = 0
  if n > 1:
    while i < n:
      tf.autograph.experimental.set_loop_options(
          shape_invariants=[(v, tf.TensorShape([None]))])
      v = tf.concat((v, [i]), 0)
      i += 1
  else:
    v = tf.constant([1, 2, 3])
  return v 

def for_with_variable_shape_growing_matrix(l):
  m = tf.constant([[0, 0], [0, 0]])
  for i in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(m, tf.TensorShape(None))])
    m = tf.pad(m, [[1, 1], [1, 1]], constant_values=i)
  return m 

def while_with_variable_shape_growing_matrix(n):
  m = tf.constant([[0, 0], [0, 0]])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(m, tf.TensorShape(None))])
    m = tf.pad(m, [[1, 1], [1, 1]], constant_values=i)
    i += 1
  return m 

def while_with_variable_shape_growing_matrix_cols(n):
  m = tf.constant([[0, 0]])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(m, tf.TensorShape([1, None]))])
    m = tf.concat((m, [[i]]), 1)
    i += 1
  return m 

def _batch_shape(self):
        return tf.TensorShape([]) 

def build(self, input_shape):
        # A small hack necessary so that train.py is completely framework-agnostic:
        input_shape = tf.TensorShape(input_shape)

        super().build(input_shape) 

def testCustomGrad(self):
    """Test for custom_grad."""
    x_shape = (tf.TensorShape([10]), tf.TensorShape([1, 10]))
    y_shape = (tf.TensorShape([]))
    dtype = np.float32
    scale1 = 5.0
    scale2 = 6.0

    def fwd(a, b):
      return tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)

    @extensions.custom_grad
    def f(a, b):
      y = fwd(a, b)

      def vjp(dy):
        return dy * scale1 * a, dy * scale2 * b

      return y, vjp

    rng = tf.random.Generator.from_seed(1234)
    x, dy = tf.nest.map_structure(lambda shape: uniform(rng, shape, dtype),
                                  [x_shape, y_shape])
    expected_y = fwd(*x)
    expected_dx = (dy * scale1 * x[0], dy * scale2 * x[1])
    y, vjp = extensions.vjp(f, *x)
    dx = vjp(dy)
    self.assertAllClose(to_tf(expected_y), to_tf(y))
    self.assertAllClose(to_tf(expected_dx), to_tf(dx)) 

def testVjp(self, has_aux):
    x_shape = (tf.TensorShape([10]), tf.TensorShape([1, 10]))
    y_shape = (tf.TensorShape([]))
    dtype = np.float32

    def f(a, b):
      y = tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)
      if has_aux:
        return y, tf_np.asarray(1)
      else:
        return y

    rng = tf.random.Generator.from_seed(1234)
    x, dy_list = tf.nest.map_structure(lambda shape: uniform(rng, shape, dtype),
                                       [x_shape, [y_shape] * 2])
    tf_x = to_tf(x)
    outputs = extensions.vjp(f, *x, has_aux=has_aux)
    if has_aux:
      y, vjp, aux = outputs
    else:
      y, vjp = outputs
    with tf.GradientTape(persistent=True) as tape:
      tape.watch(tf_x)
      outputs = f(*x)
      if has_aux:
        expected_y, expected_aux = outputs
        self.assertAllClose(to_tf(expected_aux), to_tf(aux))
      else:
        expected_y = outputs
    self.assertAllClose(to_tf(expected_y), to_tf(y))
    for dy in dy_list:
      expected_dx = tape.gradient(
          to_tf(expected_y), tf_x, output_gradients=to_tf(dy))
      self.assertAllClose(expected_dx, to_tf(vjp(dy))) 

def _replace_none_batch(x, batch_size=None):
  if batch_size is None:
    return x
  if isinstance(x, tf.Tensor) and x.shape[0] is None:
    x.set_shape([batch_size] + x.shape[1:])
    return x
  elif isinstance(x, tf.TensorShape) and x[0] is None:
    return [batch_size] + x[1:]
  return x 

def build(self, input_shape=None):
        """Build `Layer`.
        Args:
          input_shape: The shape of the input to `self.layer`.
        Raises:
          ValueError: If `Layer` does not contain a `kernel` of weights
        """

        input_shape = tf.TensorShape(input_shape).as_list()
        input_shape[0] = None
        self.input_spec = tf.keras.layers.InputSpec(shape=input_shape)

        if not self.layer.built:
            self.layer.build(input_shape)

            if not hasattr(self.layer, 'kernel'):
                raise ValueError('`WeightNorm` must wrap a layer that contains a `kernel` for weights')

            self.kernel_norm_axes = list(range(self.layer.kernel.shape.ndims))
            self.kernel_norm_axes.pop(self.filter_axis)

            self.v = self.layer.kernel

            # to avoid a duplicate `kernel` variable after `build` is called
            self.layer.kernel = None
            self.g = self.add_weight(
                name='g',
                shape=(int(self.v.shape[self.filter_axis]),),
                initializer='ones',
                dtype=self.v.dtype,
                trainable=True
            )
            self.initialized = self.add_weight(
                name='initialized',
                dtype=tf.bool,
                trainable=False
            )
            self.initialized.assign(False)

        super(WeightNorm, self).build() 

def compute_output_shape(self, input_shape):
        return tf.TensorShape(self.layer.compute_output_shape(input_shape).as_list()) 

def restore(cls, saved_model_path: str) -> "LanguageModelTF2":
        with open(saved_model_path, "rb") as fh:
            saved_data = pickle.load(fh)

        model = cls(saved_data["hyperparameters"], saved_data["vocab"])
        model.build(tf.TensorShape([None, None]))
        model.load_weights(saved_model_path)
        return model 

def _event_shape(self):
        return tf.TensorShape(self.image_shape) 

def _mvnormal_pseudo_antithetic(sample_shape,
                                mean,
                                covariance_matrix=None,
                                scale_matrix=None,
                                random_type=RandomType.PSEUDO_ANTITHETIC,
                                seed=None,
                                dtype=None):
  """Returns normal draws with the antithetic samples."""
  sample_shape = tf.TensorShape(sample_shape).as_list()
  sample_zero_dim = sample_shape[0]
  # For the antithetic sampler `sample_shape` is split evenly between
  # samples and their antithetic counterparts. In order to do the splitting
  # we expect the first dimension of `sample_shape` to be even.
  is_even_dim = tf.compat.v1.debugging.assert_equal(
      sample_zero_dim % 2,
      0,
      message='First dimension of `sample_shape` should be even for '
      'PSEUDO_ANTITHETIC random type')
  with tf.control_dependencies([is_even_dim]):
    antithetic_shape = [sample_zero_dim // 2] + sample_shape[1:]
  if random_type == RandomType.PSEUDO_ANTITHETIC:
    random_type_sample = RandomType.PSEUDO
  else:
    random_type_sample = RandomType.STATELESS
  result = _mvnormal_pseudo(
      antithetic_shape,
      mean,
      covariance_matrix=covariance_matrix,
      scale_matrix=scale_matrix,
      random_type=random_type_sample,
      seed=seed,
      dtype=dtype)
  if mean is None:
    return tf.concat([result, -result], axis=0)
  else:
    return tf.concat([result, 2 * mean - result], axis=0) 

def _prepare_index_matrix(batch_shape, num_points, dtype):
  """Prepares index matrix for index argument of `tf.gather_nd`."""
  batch_shape_reverse = batch_shape.copy()
  batch_shape_reverse.reverse()
  index_matrix = tf.constant(
      np.flip(np.transpose(np.indices(batch_shape_reverse)), -1),
      dtype=dtype)
  batch_rank = len(batch_shape)
  # Broadcast index matrix to the shape of
  # `batch_shape + [num_points] + [batch_rank]`.
  broadcasted_shape = batch_shape + [num_points] + [batch_rank]
  index_matrix = tf.expand_dims(index_matrix, -2) + tf.zeros(
      tf.TensorShape(broadcasted_shape), dtype=dtype)
  return index_matrix 

def broadcast_batch_shape(x, batch_shape):
  """Broadcasts batch shape of `x`."""
  return tf.broadcast_to(x, tf.TensorShape(batch_shape) + x.shape[-1]) 

def while_with_variable_shape_growing_vector(n):
  v = tf.constant([0, 0])
  i = 0
  while i < n:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))])
    v = tf.concat((v, [i]), 0)
    i += 1
  return v 

def for_with_variable_shape_growing_vector(l):
  v = tf.constant([0, 0])
  for i in l:
    tf.autograph.experimental.set_loop_options(
        shape_invariants=[(v, tf.TensorShape([None]))])
    v = tf.concat((v, [i]), 0)
  return v