Converts Tensor nest to a TensorShape nest, removing batch dimension.
541953
def _get_shape_without_batch_dimension(tensor_nest):
  """"""
  def _strip_batch_and_convert_to_shape(tensor):
    return tensor.get_shape()[1:]
  return nest.map_structure(_strip_batch_and_convert_to_shape, tensor_nest)
8241
def _interactive_input_tensor_to_features_dict(feature_map, hparams):
  
  inputs = tf.convert_to_tensor(feature_map["inputs"])
  input_is_image = False if len(inputs.get_shape()) < 3 else True

  x = inputs
  if input_is_image:
    x = tf.image.resize_images(x, [299, 299])
    x = tf.reshape(x, [1, 299, 299, -1])
    x = tf.to_int32(x)
  else:
    # Remove the batch dimension.
    num_samples = x[0]
    length = x[2]
    x = tf.slice(x, [3], tf.to_int32([length]))
    x = tf.reshape(x, [1, -1, 1, 1])
    # Transform into a batch of size num_samples to get that many random
    # decodes.
    x = tf.tile(x, tf.to_int32([num_samples, 1, 1, 1]))

  p_hparams = hparams.problem_hparams
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = (
      IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
  features["inputs"] = x
  return features
554637
def batch_fn(dataset, training, shapes, target_names,
             batch_size=32, eval_batch_size=32, bucket_batch_length=32,
             bucket_max_length=256, bucket_min_length=8,
             bucket_length_step=1.1, buckets=None):
  """"""
  del target_names
  # If bucketing is not specified, check if target shapes are variable.
  cur_batch_size = batch_size if training else eval_batch_size
  if buckets is None:
    variable_target_shapes = False
    target_shape = shapes[1]
    for dim in target_shape:
      if dim is None:
        variable_target_shapes = True
    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
                    "of target tensors." % variable_target_shapes)
    if variable_target_shapes:
      batch_size_per_token = cur_batch_size * bucket_batch_length
      scheme = data_reader.batching_scheme(batch_size_per_token,
                                           bucket_max_length,
                                           bucket_min_length,
                                           bucket_length_step,
                                           drop_long_sequences=training)
      buckets = (scheme["boundaries"], scheme["batch_sizes"])

  if buckets:
    tf.logging.info("Bucketing with buckets %s." % str(buckets))
    def example_length(_, target):
      return tf.shape(target)[0]
    boundaries, batch_sizes = buckets
    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
        example_length, boundaries, batch_sizes))
  else:
    dataset = dataset.padded_batch(cur_batch_size, shapes)
  return dataset
961433
def _unpack_obs(obs, space, tensorlib=tf):
    

    if (isinstance(space, gym.spaces.Dict)
            or isinstance(space, gym.spaces.Tuple)):
        prep = get_preprocessor(space)(space)
        if len(obs.shape) != 2 or obs.shape[1] != prep.shape[0]:
            raise ValueError(
                "Expected flattened obs shape of [None, {}], got {}".format(
                    prep.shape[0], obs.shape))
        assert len(prep.preprocessors) == len(space.spaces), \
            (len(prep.preprocessors) == len(space.spaces))
        offset = 0
        if isinstance(space, gym.spaces.Tuple):
            u = []
            for p, v in zip(prep.preprocessors, space.spaces):
                obs_slice = obs[:, offset:offset + p.size]
                offset += p.size
                u.append(
                    _unpack_obs(
                        tensorlib.reshape(obs_slice, [-1] + list(p.shape)),
                        v,
                        tensorlib=tensorlib))
        else:
            u = OrderedDict()
            for p, (k, v) in zip(prep.preprocessors, space.spaces.items()):
                obs_slice = obs[:, offset:offset + p.size]
                offset += p.size
                u[k] = _unpack_obs(
                    tensorlib.reshape(obs_slice, [-1] + list(p.shape)),
                    v,
                    tensorlib=tensorlib)
        return u
    else:
        return obs
120061
def export_to_tf_tensor(self, x, laid_out_x):
    
    return self.combine_slices(laid_out_x.all_slices, x.shape)
849516
def reshape_fortran(tensor, shape):
    
    return tensor.T.reshape(tuple(reversed(shape))).T
285680
def convert_batchnorm(node, **kwargs):
    
    name, input_nodes, attrs = get_inputs(node, kwargs)

    momentum = float(attrs.get("momentum", 0.9))
    eps = float(attrs.get("eps", 0.001))

    bn_node = onnx.helper.make_node(
        "BatchNormalization",
        input_nodes,
        [name],
        name=name,
        epsilon=eps,
        momentum=momentum,
        # MXNet computes mean and variance per feature for batchnorm
        # Default for onnx is across all spatial features. So disabling the parameter.
        spatial=0
    )
    return [bn_node]
1332918
def _merge_beam_dim(tensor):
  
  shape = common_layers.shape_list(tensor)
  shape[0] *= shape[1]  # batch -> batch * beam_size
  shape.pop(1)  # Remove beam dim
  return tf.reshape(tensor, shape)
560578
def nested_shape(array_or_tuple):
    
    if hasattr(array_or_tuple, 'size'):
        # pytorch tensors use V.size() to get size of tensor
        return list(array_or_tuple.size())
    elif hasattr(array_or_tuple, 'get_shape'):
        # tensorflow uses V.get_shape() to get size of tensor
        return array_or_tuple.get_shape().as_list()
    elif hasattr(array_or_tuple, 'shape'):
        return array_or_tuple.shape

    try:
        # treat object as iterable
        return [nested_shape(item) for item in list(array_or_tuple)]
    except TypeError:
        # object is not actually iterable
        # LB: Maybe we should throw an error?
        return []
234964
def convert_lstm(builder, layer, input_names, output_names, keras_layer):
    

    hidden_size = keras_layer.units
    input_size = keras_layer.input_shape[-1]
    output_all = keras_layer.return_sequences
    reverse_input = keras_layer.go_backwards

    # Keras: [W_x, W_h, b] each in I F C O
    # CoreML: I F O G; W_h and W_x are separated
    W_h, W_x, b = ([], [], [])
    keras_W_h = keras_layer.get_weights()[1].T
    W_h.append(keras_W_h[0 * hidden_size:][:hidden_size])
    W_h.append(keras_W_h[1 * hidden_size:][:hidden_size])
    W_h.append(keras_W_h[3 * hidden_size:][:hidden_size])
    W_h.append(keras_W_h[2 * hidden_size:][:hidden_size])

    keras_W_x = keras_layer.get_weights()[0].T
    W_x.append(keras_W_x[0 * hidden_size:][:hidden_size])
    W_x.append(keras_W_x[1 * hidden_size:][:hidden_size])
    W_x.append(keras_W_x[3 * hidden_size:][:hidden_size])
    W_x.append(keras_W_x[2 * hidden_size:][:hidden_size])
    if keras_layer.use_bias:
        keras_b = keras_layer.get_weights()[2]
        b.append(keras_b[0 * hidden_size:][:hidden_size])
        b.append(keras_b[1 * hidden_size:][:hidden_size])
        b.append(keras_b[3 * hidden_size:][:hidden_size])
        b.append(keras_b[2 * hidden_size:][:hidden_size])
    if len(b) == 0:
        b = None

    # Set activation type
    inner_activation_str = _get_recurrent_activation_name_from_keras(keras_layer.recurrent_activation)
    activation_str = _get_recurrent_activation_name_from_keras(keras_layer.activation)

    # Add to the network
    builder.add_unilstm(
        name = layer,
        W_h = W_h, W_x = W_x, b = b,
        hidden_size = hidden_size,
        input_size = input_size,
        input_names = input_names,
        output_names = output_names,
        inner_activation = inner_activation_str,
        cell_state_update_activation = activation_str,
        output_activation = activation_str,
        output_all = output_all,
        forget_bias = keras_layer.unit_forget_bias,
        reverse_input = reverse_input)
1372357
def convert_reshape(node, **kwargs):
    
    name, input_nodes, attrs = get_inputs(node, kwargs)

    output_shape_list = convert_string_to_list(attrs["shape"])

    initializer = kwargs["initializer"]
    output_shape_np = np.array(output_shape_list, dtype='int64')
    data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[output_shape_np.dtype]
    dims = np.shape(output_shape_np)

    output_shape_name = "reshape_attr_tensor" + str(kwargs["idx"])
    tensor_node = onnx.helper.make_tensor_value_info(output_shape_name, data_type, dims)

    initializer.append(
        onnx.helper.make_tensor(
            name=output_shape_name,
            data_type=data_type,
            dims=dims,
            vals=output_shape_list,
            raw=False,
        )
    )

    input_nodes.append(output_shape_name)

    not_supported_shape = [-2, -3, -4]

    for val in output_shape_list:
        if val in not_supported_shape:
            raise AttributeError("Reshape: Shape value not supported in ONNX", val)

    reshape_node = onnx.helper.make_node(
        "Reshape",
        input_nodes,
        [name],
        name=name
    )

    return [tensor_node, reshape_node]
671623
def _process_transform(self, data, transform, step_size):
    '''
    
    '''
    if isinstance(transform, (list,tuple,set)):
      return { t : self._transform(data,t,step_size) for t in transform }
    elif isinstance(transform, dict):
      return { tn : self._transform(data,tf,step_size) for tn,tf in transform.items() }
    return self._transform(data,transform,step_size)
43100
def split_last_dimension(x, n):
  
  x_shape = common_layers.shape_list(x)
  m = x_shape[-1]
  if isinstance(m, int) and isinstance(n, int):
    assert m % n == 0
  return tf.reshape(x, x_shape[:-1] + [n, m // n])
1048714
def _batch(self,
               vectors,
               batch_size,
               num,
               show_progressbar,
               return_names):
        """"""
        vectors = self.normalize(vectors)

        # Single transpose, makes things faster.
        reference_transposed = self.norm_vectors.T

        for i in tqdm(range(0, len(vectors), batch_size),
                      disable=not show_progressbar):

            distances = vectors[i: i+batch_size].dot(reference_transposed)
            # For safety we clip
            distances = np.clip(distances, a_min=.0, a_max=1.0)
            if num == 1:
                sorted_indices = np.argmax(distances, 1)[:, None]
            else:
                sorted_indices = np.argpartition(-distances, kth=num, axis=1)
                sorted_indices = sorted_indices[:, :num]
            for lidx, indices in enumerate(sorted_indices):
                dists = distances[lidx, indices]
                if return_names:
                    dindex = np.argsort(-dists)
                    yield [(self.indices[indices[d]], dists[d])
                           for d in dindex]
                else:
                    yield list(-1 * np.sort(-dists))
897784
public static SparseTensorBuilder copyOf(TensorBase tensor) {
    SparseTensorBuilder builder = new SparseTensorBuilder(tensor.getDimensionNumbers(), 
        tensor.getDimensionSizes());
    Iterator<KeyValue> initialWeightIter = tensor.keyValueIterator();
    while (initialWeightIter.hasNext()) {
      KeyValue keyValue = initialWeightIter.next();
      builder.put(keyValue.getKey(), keyValue.getValue());
    }
    return builder;
  }
634109
def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
    """"""
    # Store the original shape of the mean and variance.
    mean_shape = mean.get_shape()
    variance_shape = variance.get_shape()
    # The fused batch norm expects the mean, variance, gamma and beta
    # tensors to have dimension 1, so we flatten them to remove the
    # extra dimensions. In addition, it expects the input_batch to have
    # dimension 4, so we reshape it accordingly.
    gamma_flatten = tf.reshape(self._gamma, shape=(self._num_channels,))
    beta_flatten = tf.reshape(self._beta, shape=(self._num_channels,))
    flatten_mean = tf.reshape(mean, shape=(self._num_channels,))
    flatten_variance = tf.reshape(variance, shape=(self._num_channels,))
    use_batch_stats = tf.convert_to_tensor(use_batch_stats)

    input_shape = input_batch.get_shape()
    output_shape = [-1] + input_shape.as_list()[1:]

    flat_image_size = np.prod(self._image_shape, dtype=np.int32)

    if len(self._data_format) == 4:
      fusable_data_format = self._data_format
      fusable_batch = input_batch
    elif self._channel_index == 1 and self._image_shape:
      fusable_data_format = "NCHW"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, self._num_channels, 1, flat_image_size))
    else:
      # The CPU implementation of FusedBatchNorm only supports NHWC tensor
      # format for now.
      fusable_data_format = "NHWC"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, 1, flat_image_size, self._num_channels))

    common_args = {
        "scale": gamma_flatten,
        "offset": beta_flatten,
        "epsilon": self._eps,
        "data_format": fusable_data_format,
        "name": "batch_norm"
    }

    def use_batch_stats_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=None,
          variance=None,
          is_training=True,
          **common_args)

    def moving_average_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=flatten_mean,
          variance=flatten_variance,
          is_training=False,
          **common_args)

    batch_norm_op, mean, variance = utils.smart_cond(
        use_batch_stats, use_batch_stats_fused_batch_norm,
        moving_average_fused_batch_norm)

    if len(self._data_format) != 4:
      batch_norm_op = tf.reshape(batch_norm_op, output_shape)
    mean = tf.reshape(mean, mean_shape)
    variance = tf.reshape(variance, variance_shape)
    return batch_norm_op, mean, variance
341544
def transpose(self, permutation: Optional[List[int]] = None) -> 'TensorFluent':
        '''
        '''
        if permutation == []:
            return self
        t = tf.transpose(self.tensor, permutation) if permutation != [] else self.tensor
        scope = self.scope.as_list()
        batch = self.batch
        return TensorFluent(t, scope, batch=batch)