Python tensorflow.compat.v2.TensorSpec() Examples

def __init__(self, input_tensor_spec=None, state_spec=(), name=None):
    """Creates an instance of `Network`.

    Args:
      input_tensor_spec: A nest of `tf.TypeSpec` representing the
        input observations.  Optional.  If not provided, `create_variables()`
        will fail unless a spec is provided.
      state_spec: A nest of `tensor_spec.TensorSpec` representing the state
        needed by the network. Default is `()`, which means no state.
      name: (Optional.) A string representing the name of the network.
    """
    # Disable autocast because it may convert bfloats to other types, breaking
    # our spec checks.
    super(Network, self).__init__(name=name, autocast=False)
    common.check_tf1_allowed()

    # Required for summary() to work.
    self._is_graph_network = False

    self._input_tensor_spec = input_tensor_spec
    # NOTE(ebrevdo): Would have preferred to call this output_tensor_spec, but
    # looks like keras.Layer already reserves that one.
    self._network_output_spec = None
    self._state_spec = state_spec 

def __init__(self, state_space_size, unroll_length=1):
    self._state_space_size = state_space_size
    # Creates simple dynamics (T stands for transition):
    #   states = [0, 1, ... len(state_space_size - 1)] + [STOP]
    #   actions = [-1, 1]
    #   T(s, a) = s + a  iff (s + a) is a valid state
    #           = STOP   otherwise
    self._action_space = [-1, 1]
    self._current_state = None
    self._env_spec = common.EnvOutput(
        reward=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.float32),
        done=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.bool),
        observation={
            'f1':
                tf.TensorSpec(
                    shape=[unroll_length + 1, 4, 10], dtype=tf.float32),
            'f2':
                tf.TensorSpec(
                    shape=[unroll_length + 1, 7, 10, 2], dtype=tf.float32)
        },
        info=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.string)) 

def test_run_eval_actor_once(self):
    hparams = {}
    hparams['max_iter'] = 1
    hparams['num_episodes_per_iter'] = 5
    hparams['logdir'] = os.path.join(FLAGS.test_tmpdir, 'model')

    mock_problem = testing_utils.MockProblem(unroll_length=FLAGS.unroll_length)
    agent = mock_problem.get_agent()
    ckpt_manager = _get_ckpt_manager(hparams['logdir'], agent=agent)
    ckpt_manager.save(checkpoint_number=0)

    # Create a no-op gRPC server that responds to Aggregator RPCs.
    server_address = 'unix:/tmp/eval_actor_test_grpc'
    server = grpc.Server([server_address])

    @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.string)])
    def eval_enqueue(_):
      return []

    server.bind(eval_enqueue, batched=False)

    server.start()

    eval_actor.run_with_aggregator(mock_problem, server_address, hparams) 

def testReadWriteSpecs(self):
    logdir = FLAGS.test_tmpdir
    specs = {
        'a': tf.TensorSpec(shape=(2, 3), dtype=tf.float32),
        'b': {
            'b_1': tf.TensorSpec(shape=(5,), dtype=tf.string),
            'b_2': tf.TensorSpec(shape=(5, 6), dtype=tf.int32),
        }
    }
    utils.write_specs(logdir, specs)
    # Now read and verify
    specs_read = utils.read_specs(logdir)

    def _check_equal(sp1, sp2):
      self.assertEqual(sp1, sp2)

    tf.nest.map_structure(_check_equal, specs, specs_read) 

def test_model_id_updater(self):
    loc = 0.0
    scale = 0.5
    action_spec = tensor_spec.BoundedTensorSpec([1], tf.float32, tf.float32.min,
                                                tf.float32.max)
    wrapped_policy = DistributionPolicy(
        distribution=tfp.distributions.Normal([loc], [scale]),
        time_step_spec=self._time_step_spec,
        action_spec=action_spec,
        info_spec={
            'test_info':
                tf.TensorSpec(shape=(1,), dtype=tf.int32, name='test_info')
        })
    updater_info_spec = {
        'model_id': tf.TensorSpec(shape=(1,), dtype=tf.int32, name='model_id')
    }
    updater_info_spec.update(wrapped_policy.info_spec)
    policy = policy_info_updater_wrapper.PolicyInfoUpdaterWrapper(
        policy=wrapped_policy,
        info_spec=updater_info_spec,
        updater_fn=ModelIdUpdater(),
        name='model_id_updater')

    self.assertEqual(policy.time_step_spec, self._time_step_spec)
    self.assertEqual(policy.action_spec, action_spec)

    observations = tf.constant([[1, 2], [3, 4]], dtype=tf.float32)
    time_step = ts.restart(observations, batch_size=2)
    action_step = policy.action(time_step)
    distribution_step = policy.distribution(time_step)

    tf.nest.assert_same_structure(action_spec, action_step.action)
    tf.nest.assert_same_structure(action_spec, distribution_step.action)

    self.assertListEqual(list(self.evaluate(action_step.info['model_id'])), [2])
    self.assertListEqual(
        list(self.evaluate(distribution_step.info['model_id'])), [2]) 

def _write_tensor_specs(initial_agent_state: Any,
                        env_output: common.EnvOutput,
                        agent_output: common.AgentOutput,
                        actor_action: common.ActorAction,
                        loss_type: Optional[int] = common.AC_LOSS):
  """Writes tensor specs of ActorOutput tuple to disk.

  Args:
    initial_agent_state: A tensor or nested structure of tensor without any time
      or batch dimensions.
    env_output: An instance of `EnvOutput` where individual tensors don't have
      time and batch dimensions.
    agent_output: An instance of `AgentOutput` where individual tensors don't
      have time and batch dimensions.
    actor_action: An instance of `ActorAction`.
    loss_type: A scalar int denoting the loss type.
  """
  actor_output = common.ActorOutput(
      initial_agent_state,
      env_output,
      agent_output,
      actor_action,
      loss_type,
      info='')
  specs = tf.nest.map_structure(tf.convert_to_tensor, actor_output)
  specs = tf.nest.map_structure(tf.TensorSpec.from_tensor, specs)
  env_output = tf.nest.map_structure(add_time_dimension, specs.env_output)
  agent_output = tf.nest.map_structure(add_time_dimension, specs.agent_output)
  actor_action = tf.nest.map_structure(add_time_dimension, specs.actor_action)
  specs = specs._replace(
      env_output=env_output,
      agent_output=agent_output,
      actor_action=actor_action)
  utils.write_specs(FLAGS.logdir, specs) 

def add_time_dimension(s: tf.TensorSpec):
  return tf.TensorSpec([FLAGS.unroll_length + 1] + s.shape.as_list(), s.dtype) 

def __init__(self, unroll_length=1):
    super(MockAgent, self).__init__(name='mock_agent')
    self._state_size = 5
    # This matches the action space of MockEnv
    self._action_space_size = 2
    self._logits_layer = tf.keras.layers.Dense(self._action_space_size)
    self._agent_spec = common.AgentOutput(
        policy_logits=tf.TensorSpec(
            shape=[unroll_length + 1, 2], dtype=tf.float32),
        baseline=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.float32)) 

def assert_matches_spec(specs, tensor_list):
  """Assert that a list of tensors matches the given TensorSpecs."""
  # Weirdly `tf.nest.pack_sequence_as` doesn't fail if tensor_list doesn't
  # conform to the specs type. So first pack the sequence, then explicitly
  # check the compatibility of each tensor with the corresponding spec.
  packed_tensors = tf.nest.pack_sequence_as(specs, tensor_list)
  packed_tensors = tf.nest.map_structure(tf.convert_to_tensor, packed_tensors)

  def is_compatible(sp, tensor):
    assert sp.is_compatible_with(
        tensor), 'TensorSpec {} is not compatible with tensor {}'.format(
            sp, tensor)

  tf.nest.map_structure(is_compatible, specs, packed_tensors) 

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 setUp(self):
    super(PolicyInfoUpdaterWrapperTest, self).setUp()
    self._obs_spec = tensor_spec.TensorSpec([2], tf.float32)
    self._time_step_spec = ts.time_step_spec(self._obs_spec) 

def testIncompatibleStructureInputs(self):
    with self.assertRaisesRegex(
        ValueError,
        r'`nested_layers` and `input_spec` do not have matching structures'):
      nest_map.NestMap(
          tf.keras.layers.Dense(8),
          input_spec={'ick': tf.TensorSpec(8, tf.float32)})

    with self.assertRaisesRegex(
        ValueError,
        r'`inputs` and `self.nested_layers` do not have matching structures'):
      net = nest_map.NestMap(tf.keras.layers.Dense(8))
      net.create_variables({'ick': tf.TensorSpec((1,), dtype=tf.float32)})

    with self.assertRaisesRegex(
        ValueError,
        r'`inputs` and `self.nested_layers` do not have matching structures'):
      net = nest_map.NestMap(tf.keras.layers.Dense(8))
      net({'ick': tf.constant([[1.0]])})

    with self.assertRaisesRegex(
        ValueError,
        r'`network_state` and `state_spec` do not have matching structures'):
      net = nest_map.NestMap(
          tf.keras.layers.LSTM(8, return_state=True, return_sequences=True))
      net(tf.ones((1, 2)), network_state=(tf.ones((1, 1)), ())) 

def __init__(self, time_step_spec, net):
    super(MyPolicy, self).__init__(
        time_step_spec,
        action_spec=tf.TensorSpec((None,), tf.float32))
    self._net = net 

def get_state_spec(layer: tf.keras.layers.Layer) -> types.NestedTensorSpec:
  """Extracts the state spec from a layer.

  Args:
    layer: The layer to extract from; can be a `Network`.

  Returns:
    The state spec.

  Raises:
    TypeError: If `layer` is a subclass of `tf.keras.layers.RNN` (it must
      be wrapped by an `RNNWrapper` object).
    ValueError: If `layer` is a Keras layer and `create_variables` has
      not been called on it.
  """
  if isinstance(layer, Network):
    return layer.state_spec

  if isinstance(layer, tf.keras.layers.RNN):
    raise TypeError("RNN Layer must be wrapped inside "
                    "`tf_agents.keras_layers.RNNWrapper`: {}".format(layer))

  initial_state = getattr(layer, "get_initial_state", None)
  state_size = getattr(layer, "state_size", None)
  if initial_state is not None and state_size is None:
    raise ValueError(
        "Layer lacks a `state_size` property.  Unable to extract state "
        "spec: {}".format(layer))
  state_spec = ()
  if state_size is not None:
    state_spec = tf.nest.map_structure(
        lambda s: tf.TensorSpec(dtype=layer.dtype, shape=s), state_size)

  return state_spec 

def export_serving_model(tf_transform_output, model, output_dir):
  """Exports a keras model for serving.

  Args:
    tf_transform_output: Wrapper around output of tf.Transform.
    model: A keras model to export for serving.
    output_dir: A directory where the model will be exported to.
  """
  # The layer has to be saved to the model for keras tracking purpases.
  model.tft_layer = tf_transform_output.transform_features_layer()

  @tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    """Serving tf.function model wrapper."""
    feature_spec = RAW_DATA_FEATURE_SPEC.copy()
    feature_spec.pop(LABEL_KEY)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
    transformed_features = model.tft_layer(parsed_features)
    outputs = model(transformed_features)
    classes_names = tf.constant([['0', '1']])
    classes = tf.tile(classes_names, [tf.shape(outputs)[0], 1])
    return {'classes': classes, 'scores': outputs}

  concrete_serving_fn = serve_tf_examples_fn.get_concrete_function(
      tf.TensorSpec(shape=[None], dtype=tf.string, name='inputs'))
  signatures = {'serving_default': concrete_serving_fn}

  # This is required in order to make this model servable with model_server.
  versioned_output_dir = os.path.join(output_dir, '1')
  model.save(versioned_output_dir, save_format='tf', signatures=signatures) 

def eval_on_shapes(f, static_argnums=()):
  """Returns a function that evaluates `f` given input shapes and dtypes.

  It transforms function `f` to a function that performs the same computation as
  `f` but only on shapes and dtypes (a.k.a. shape inference).

  Args:
    f: the function to be transformed.
    static_argnums: See documentation of `jit`.

  Returns:
    A function whose input arguments can be either the same as `f`'s or only
    their shapes/dtypes represented by `TensorSpec`, and whose return values are
    `TensorSpec`s with the same nested structure as `f`'s return values.
  """
  # TODO(wangpeng): tf.function could add a knob to turn off materializing the
  #   graph, so that we don't waste computation and memory when we just want
  #   shape inference.
  tf_f = jit(f, static_argnums=static_argnums).tf_function

  # pylint: disable=missing-docstring
  def f_return(*args):

    def abstractify(x):
      x = _canonicalize_jit_arg(x)
      if isinstance(x, (tf.Tensor, tf_np.ndarray)):
        return tf.TensorSpec(x.shape, x.dtype)
      else:
        return x

    def to_tensor_spec(x):
      if isinstance(x, tf.Tensor):
        return tf.TensorSpec(x.shape, x.dtype)
      else:
        return x

    new_args = []
    for i, arg in enumerate(args):
      if i in static_argnums:
        new_args.append(arg)
      else:
        new_args.append(tf.nest.map_structure(abstractify, arg))
    res = tf_f.get_concrete_function(*new_args).structured_outputs

    return tf.nest.map_structure(to_tensor_spec, res)

  # Provides access to `tf_f` for testing purpose.
  f_return._tf_function = tf_f  # pylint: disable=protected-access
  return f_return 

def testCreateAndCall(self):
    net = sequential.Sequential([
        nest_map.NestMap(
            {'inp1': tf.keras.layers.Dense(8),
             'inp2': sequential.Sequential([
                 tf.keras.layers.Conv2D(2, 3),
                 # Convert 3 inner dimensions to [8] for RNN.
                 inner_reshape.InnerReshape([None] * 3, [8]),
             ]),
             'inp3': tf.keras.layers.LSTM(
                 8, return_state=True, return_sequences=True)}),
        nest_map.NestFlatten(),
        tf.keras.layers.Add()])
    self.assertEqual(
        net.state_spec,
        ({
            'inp1': (),
            'inp2': (),
            'inp3': (2 * [tf.TensorSpec(shape=(8,), dtype=tf.float32)],),
        },))
    output_spec = net.create_variables(
        {
            'inp1': tf.TensorSpec(shape=(3,), dtype=tf.float32),
            'inp2': tf.TensorSpec(shape=(4, 4, 2,), dtype=tf.float32),
            'inp3': tf.TensorSpec(shape=(3,), dtype=tf.float32),
        })
    self.assertEqual(output_spec, tf.TensorSpec(shape=(8,), dtype=tf.float32))

    inputs = {
        'inp1': tf.ones((8, 10, 3), dtype=tf.float32),
        'inp2': tf.ones((8, 10, 4, 4, 2), dtype=tf.float32),
        'inp3': tf.ones((8, 10, 3), dtype=tf.float32)
    }
    output, next_state = net(inputs)
    self.assertEqual(output.shape, tf.TensorShape([8, 10, 8]))
    self.assertEqual(
        tf.nest.map_structure(lambda t: t.shape, next_state),
        ({
            'inp1': (),
            'inp2': (),
            'inp3': (2 * [tf.TensorShape([8, 8])],),
        },))

    # Test passing in a state.
    output, next_state = net(inputs, next_state)
    self.assertEqual(output.shape, tf.TensorShape([8, 10, 8]))
    self.assertEqual(
        tf.nest.map_structure(lambda t: t.shape, next_state),
        ({
            'inp1': (),
            'inp2': (),
            'inp3': (2 * [tf.TensorShape([8, 8])],),
        },)) 

def tf_abstract_eval(f):
  """Returns a function that evaluates `f` given input shapes and dtypes.

  It transforms function `f` to a function that performs the same computation as
  `f` but only on shapes and dtypes (a.k.a. shape inference).

  Args:
    f: the function to be transformed.

  Returns:
    A function whose input arguments can be either the same as `f`'s or only
    their shapes/dtypes represented by `ShapeDtype`, and whose return values are
    `ShapeDtype`s with the same nested structure as `f`'s return values.
  """
  f_shape = tf_np_extensions.eval_on_shapes(f)
  def from_shape_type(x):
    if isinstance(x, ShapeDtype):
      return tf.TensorSpec(x.shape, x.dtype)
    else:
      return x
  def to_shape_type(x):  # pylint: disable=missing-docstring
    # TODO(wangpeng): handle partial output shapes using `tf.shape`.
    def to_numpy_shape(s):
      if s.is_fully_defined():
        return tuple(s.as_list())
      else:
        raise ValueError("The output shapes (%s) of the dry-run'ed function are"
                         ' not fully defined.' % s)
    def to_numpy_dtype(t):
      return np.dtype(t.as_numpy_dtype)
    if isinstance(x, tf.TensorSpec):
      return ShapeDtype(to_numpy_shape(x.shape), to_numpy_dtype(x.dtype))
    else:
      return x
  def f_return(*args):
    args = tf.nest.map_structure(from_shape_type, args)
    res = f_shape(*args)
    return tf.nest.map_structure(to_shape_type, res)
  return f_return


# The arguments order is different from tf_np_extensions.uniform