<commit_before># Copyright 2019 DeepMind Technologies Limited. 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.
# ==============================================================================
"""Basic Haiku modules and functions."""

import functools
import types
from typing import Any, Callable, Iterable, Optional, Type

from haiku._src import base
from haiku._src import initializers
from haiku._src import module
from haiku._src.typing import PRNGKey
import jax
from jax import lax
import jax.numpy as jnp
import numpy as np

# If you are forking replace this block with `import haiku as hk`.
hk = types.ModuleType("haiku")
hk.get_parameter = base.get_parameter
hk.initializers = initializers
hk.Module = module.Module
del base, module, initializers


# Utility and activation functions.
def one_hot(x, num_classes, dtype=jnp.float32):
  """Returns a one-hot version of indices.

  DEPRECATED: Use ``jax.nn.one_hot(x, num_classes).astype(dtype)`` instead.

  Args:
    x: A tensor of indices.
    num_classes: Number of classes in the one-hot dimension.
    dtype: The dtype.

  Returns:
    The one-hot tensor. If indices' shape is [A, B, ...], shape is
      [A, B, ... num_classes].
  """
  return jax.nn.one_hot(x, num_classes).astype(dtype)


def multinomial(rng, logits, num_samples):
  """Draws samples from a multinomial distribution.

  Args:
    rng: A JAX PRNGKey.
    logits: Unnormalized log-probabilities, of shape
      ``[batch_size, categories]`` or ``[categories]``.
    num_samples: Number of samples to draw.

  Returns:
    Chosen categories, of shape ``[batch_size, num_samples]`` or
    ``[num_samples]``.
  """
  # NOTE(tycai): Currently, tf.multinomial uses CDF for non-XLA CPU only.
  # We may want to switch to the Gumbel trick as used in XLA.
  if len(logits.shape) > 2 or not logits.shape:
    raise ValueError("Logits must be rank-1 or rank-2.")
  probs = jax.nn.softmax(logits)
  probs = jnp.cumsum(probs, axis=-1)
  # Special-case num_samples == 1 due to TPU padding, as in TF2XLA.
  # https://github.com/tensorflow/tensorflow/blob/b1608511d5a50d05825c4025b0c347e8689a241f/tensorflow/compiler/tf2xla/kernels/categorical_op.cc#L79
  if num_samples == 1:
    a = jax.random.uniform(rng, logits.shape[:-1] + (1,))
    out = jnp.argmin(a > probs, axis=-1)
    return out[..., None]
  else:
    a = jax.random.uniform(rng, (num_samples,) + logits.shape[:-1] + (1,))
    out = jnp.argmin(a > probs, axis=-1)
    return jnp.transpose(out)


# Common modules.
class Sequential(hk.Module):
  """Sequentially calls the given list of layers.

  Note that :class:`Sequential` is limited in the range of possible
  architectures it can handle. This is a deliberate design decision;
  :class:`Sequential` is only meant to be used for the simple case of fusing
  together modules/ops where the input of a particular module/op is the output
  of the previous one.

  Another restriction is that it is not possible to have extra arguments in the
  :meth:`__call__` method that are passed to the constituents of the module -
  for example, if there is a :class:`BatchNorm` module in :class:`Sequential`
  and the user wishes to switch the ``is_training`` flag. If this is the desired
  use case, the recommended solution is to subclass :class:`Module` and
  implement ``__call__``:

      >>> class CustomModule(hk.Module):<commit_msg>[hk] Switch hk.multinomial to use jax.random.categorical<commit_after># Copyright 2019 DeepMind Technologies Limited. 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.
# ==============================================================================
"""Basic Haiku modules and functions."""

import functools
import types
from typing import Any, Callable, Iterable, Optional, Type

from haiku._src import base
from haiku._src import initializers
from haiku._src import module
from haiku._src.typing import PRNGKey
import jax
from jax import lax
import jax.numpy as jnp
import numpy as np

# If you are forking replace this block with `import haiku as hk`.
hk = types.ModuleType("haiku")
hk.get_parameter = base.get_parameter
hk.initializers = initializers
hk.Module = module.Module
del base, module, initializers


# Utility and activation functions.
def one_hot(x, num_classes, dtype=jnp.float32):
  """Returns a one-hot version of indices.

  DEPRECATED: Use ``jax.nn.one_hot(x, num_classes).astype(dtype)`` instead.

  Args:
    x: A tensor of indices.
    num_classes: Number of classes in the one-hot dimension.
    dtype: The dtype.

  Returns:
    The one-hot tensor. If indices' shape is [A, B, ...], shape is
      [A, B, ... num_classes].
  """
  return jax.nn.one_hot(x, num_classes).astype(dtype)


def multinomial(rng, logits, num_samples):
  """Draws samples from a multinomial distribution.

  DEPRECATED: Use ``jax.random.categorical`` instead.

  Args:
    rng: A JAX PRNGKey.
    logits: Unnormalized log-probabilities, where last dimension is categories.
    num_samples: Number of samples to draw.

  Returns:
    Chosen categories, of shape ``logits.shape[:-1] + (num_samples,)``.
  """
  if num_samples != 1:
    shape = (num_samples,) + logits.shape[:-1]
  else:
    shape = None

  samples = jax.random.categorical(rng, logits, shape=shape)

  # Return expected shape.
  if num_samples != 1:
    return jnp.moveaxis(samples, 0, -1)
  else:
    return samples[..., None]


# Common modules.
class Sequential(hk.Module):
  """Sequentially calls the given list of layers.

  Note that :class:`Sequential` is limited in the range of possible
  architectures it can handle. This is a deliberate design decision;
  :class:`Sequential` is only meant to be used for the simple case of fusing
  together modules/ops where the input of a particular module/op is the output
  of the previous one.

  Another restriction is that it is not possible to have extra arguments in the
  :meth:`__call__` method that are passed to the constituents of the module -
  for example, if there is a :class:`BatchNorm` module in :class:`Sequential`
  and the user wishes to switch the ``is_training`` flag. If this is the desired
  use case, the recommended solution is to subclass :class:`Module` and
  implement ``__call__``:

      >>> class CustomModule(hk.Module):
{'fraction_non_alphanumeric': 0.1142502756339581, 'fraction_numerical': 0.009922822491730982, 'mean_word_length': 4.0325936199722605, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 12, 'https://': 1, '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': '3262976', 'n_tokens_mistral': 2248, 'n_tokens_neox': 2144, 'n_words': 857}