<commit_before># Copyright 2019 The Sonnet Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Reshaping Sonnet modules."""

from __future__ import absolute_import
from __future__ import division
# from __future__ import google_type_annotations
from __future__ import print_function

import numpy as np
from sonnet.src import base
from sonnet.src import once
import tensorflow as tf


def reshape(inputs, output_shape, preserve_dims=1, name="reshape"):
  """Reshapes input Tensor, preserving the batch dimension.

      >>> snt.reshape(
      ...   tf.ones([10, 2, 2, 3]), output_shape=[-1], preserve_dims=1)
      <tf.Tensor: ... shape=(10, 12), ...>

      >>> snt.reshape(
      ...   tf.ones([10, 2, 2, 3]), output_shape=[-1], preserve_dims=2)
      <tf.Tensor: ... shape=(10, 2, 6), ...>

  Args:
    inputs: A tensor of shape
      ``[b_1, b_2, ..., b_preserve_dims, b_preserve_dims+1, ...]``.
    output_shape: Shape to reshape the input tensor to while preserving its
      first ``preserve_dims`` dimensions. When the special value -1 appears in
      ``shape`` the corresponding size is automatically inferred. Note that -1
      can only appear once in ``shape``. To flatten all non-batch dimensions,
      :func:`flatten` can also be used.
    preserve_dims: Number of leading dimensions that will not be reshaped. For
      example, given an input tensor with shape ``[B, H, W, C, D]``, and
      argument ``shape`` equal to ``(-1, D)``:

      * ``preserve_dims=1`` returns a tensor with shape ``[B, H*W*C, D]``.
      * ``preserve_dims=2`` returns a tensor with shape ``[B, H, W*C, D]``.
      * ``preserve_dims=3`` returns a tensor with shape ``[B, H, W, C, D]``.
      * ``preserve_dims=4`` returns a tensor with shape
        ``[B, H, W, C, 1, D]``.
      * ``preserve_dims>=5`` will throw an error on build unless ``D=1``.

      The preserved dimensions can be unknown at building time.
    name: Optional name scope used to help debugging.

  Returns:
    A tensor of shape
      ``[b_1, b_2, ..., b_preserve_dims, b_reshape_1, b_reshape_2, ...]``,
      with reshaping defined by the constructor ``shape`` parameter.

  Raises:
    ValueError: If ``output_shape`` is incompatible with shape of the
      ``inputs``; or if ``output_shape`` contains more than one wildcard -1;
      or if the ``inputs`` rank is less than ``preserved_dims``; or if
      the ``inputs`` shape contains unknown, non-preserved dimensions
      (except when the unknown dimension is the only non-preserved
      dimension and doesn't actually need reshaping).
  """
  if preserve_dims <= 0:
    raise ValueError("Argument preserve_dims should be >= 1.")

  with tf.name_scope(name):
    return _batch_reshape(inputs, output_shape, preserve_dims)


def flatten(inputs, name="flatten"):
  """Flattens the input tensor preserving only the batch dimension.

      >>> snt.flatten(tf.ones([5, 3, 2, 2]))
      <tf.Tensor: ... shape=(5, 12), ...>

  Args:
    inputs: An input tensor to flatten.
    name: Optional name scope used to help debugging.

  Returns:
    A tensor corresponding to the input tensor but with all non-batch
      dimensions flattened.
  """
  with tf.name_scope(name):
    return _batch_reshape(inputs, output_shape=(-1,), preserve_dims=1)


def _extract_input_shape(inputs, preserve_dims=1) -> tf.TensorShape:
  """Extracts the shape minus ignored dimensions from ``inputs``.

  ::
      >>> _extract_input_shape(tf.ones([1, 2, 3]))
      TensorShape([2, 3])<commit_msg>Changed reshape and flatten to call to Reshape/Flatten directly<commit_after># Copyright 2019 The Sonnet Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Reshaping Sonnet modules."""

from __future__ import absolute_import
from __future__ import division
# from __future__ import google_type_annotations
from __future__ import print_function

import numpy as np
from sonnet.src import base
from sonnet.src import once
import tensorflow as tf


def reshape(inputs, output_shape, preserve_dims=1, name=None):
  """A shortcut for applying :class:`Reshape` to the ``inputs``."""
  return Reshape(output_shape, preserve_dims, name=name)(inputs)


def flatten(inputs, name="flatten"):
  """A shortcut for applying :class:`Flatten` to the ``inputs``."""
  return Flatten(name=name)(inputs)


def _extract_input_shape(inputs, preserve_dims=1) -> tf.TensorShape:
  """Extracts the shape minus ignored dimensions from ``inputs``.

  ::
      >>> _extract_input_shape(tf.ones([1, 2, 3]))
      TensorShape([2, 3])
{'fraction_non_alphanumeric': 0.13770147161878066, 'fraction_numerical': 0.014190609670637701, 'mean_word_length': 3.904639175257732, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1165154', 'n_tokens_mistral': 1706, 'n_tokens_neox': 1637, 'n_words': 643}