<commit_before># Copyright 2020 Google LLC
#
# 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
#
#     https://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.

__all__ = ['BaseVar', 'BaseState', 'RandomState', 'TrainRef', 'StateVar', 'TrainVar', 'VarCollection']

import abc
from contextlib import contextmanager
from typing import List, Union, Tuple, Optional, Iterable, Dict, Iterator, Callable

import jax
import jax.numpy as jn
import jax.random as jr
import numpy as np

from objax.typing import JaxArray
from objax.util import map_to_device, Renamer


def reduce_mean(x: JaxArray) -> JaxArray:
    return x.mean(0)


class BaseVar(abc.ABC):
    """The abstract base class to represent objax variables."""

    def __init__(self, reduce: Optional[Callable[[JaxArray], JaxArray]]):
        """Constructor for BaseVar class.

        Args:
            reduce: a function that takes an array of shape ``(n, *dims)`` and returns one of shape ``(*dims)``. Used to
                    combine the multiple states produced in an objax.Vectorize or an objax.Parallel call.
        """
        self._reduce = reduce

    @property
    @abc.abstractmethod
    def value(self) -> JaxArray:
        raise NotImplementedError('Pure method')

    @value.setter
    @abc.abstractmethod
    def value(self, tensor: JaxArray):
        raise NotImplementedError('Pure method')

    def assign(self, tensor: JaxArray, check=True):
        """Sets the value of the variable."""
        if check:
            self.assert_assigned_type_and_shape_match(tensor)
        self.value = tensor

    def reduce(self, tensors: JaxArray):
        """Method called by Parallel and Vectorize to reduce a multiple-device (or batched in case of vectoriaation)
        value to a single device."""
        if self._reduce:
            self.assign(self._reduce(tensors))

    def assert_assigned_type_and_shape_match(self, tensor):
        assert isinstance(tensor, JaxArray.__args__), \
            f'Assignments to variable must be an instance of JaxArray,  but received f{type(tensor)}.'

        def shape_and_device(array):
            if isinstance(array, jax.interpreters.pxla.ShardedDeviceArray):
                return array.shape[0], array.shape[1:]
            else:
                return None, array.shape
        tensor_device, tensor_shape = shape_and_device(tensor)
        self_device, self_shape = shape_and_device(self.value)

        device_mismatch_error = f"Can not replicate a variable that is currently on " \
                                f"{self_device} devices to {tensor_device} devices."
        assert (tensor_device is None) or (self_device is None) or (self_device == tensor_device), device_mismatch_error

        tracer_types = (jax.interpreters.partial_eval.JaxprTracer,
                        jax.interpreters.partial_eval.DynamicJaxprTracer)

        shorter_length = min(len(tensor.shape), len(self.value.shape))
        is_special_ok = (isinstance(tensor, tracer_types) or isinstance(self.value, tracer_types)) and \
                        self.value.shape[-shorter_length:] == tensor.shape[-shorter_length:]

        shape_mismatch_error = f"Assign can not change shape of variable. The current variable shape is {self_shape}," \
                               f" but the requested new shape is {tensor_shape}."<commit_msg>Fix linting error in assign reshape assert<commit_after># Copyright 2020 Google LLC
#
# 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
#
#     https://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.

__all__ = ['BaseVar', 'BaseState', 'RandomState', 'TrainRef', 'StateVar', 'TrainVar', 'VarCollection']

import abc
from contextlib import contextmanager
from typing import List, Union, Tuple, Optional, Iterable, Dict, Iterator, Callable

import jax
import jax.numpy as jn
import jax.random as jr
import numpy as np

from objax.typing import JaxArray
from objax.util import map_to_device, Renamer


def reduce_mean(x: JaxArray) -> JaxArray:
    return x.mean(0)


class BaseVar(abc.ABC):
    """The abstract base class to represent objax variables."""

    def __init__(self, reduce: Optional[Callable[[JaxArray], JaxArray]]):
        """Constructor for BaseVar class.

        Args:
            reduce: a function that takes an array of shape ``(n, *dims)`` and returns one of shape ``(*dims)``. Used to
                    combine the multiple states produced in an objax.Vectorize or an objax.Parallel call.
        """
        self._reduce = reduce

    @property
    @abc.abstractmethod
    def value(self) -> JaxArray:
        raise NotImplementedError('Pure method')

    @value.setter
    @abc.abstractmethod
    def value(self, tensor: JaxArray):
        raise NotImplementedError('Pure method')

    def assign(self, tensor: JaxArray, check=True):
        """Sets the value of the variable."""
        if check:
            self.assert_assigned_type_and_shape_match(tensor)
        self.value = tensor

    def reduce(self, tensors: JaxArray):
        """Method called by Parallel and Vectorize to reduce a multiple-device (or batched in case of vectoriaation)
        value to a single device."""
        if self._reduce:
            self.assign(self._reduce(tensors))

    def assert_assigned_type_and_shape_match(self, tensor):
        assert isinstance(tensor, JaxArray.__args__), \
            f'Assignments to variable must be an instance of JaxArray,  but received f{type(tensor)}.'

        def shape_and_device(array):
            if isinstance(array, jax.interpreters.pxla.ShardedDeviceArray):
                return array.shape[0], array.shape[1:]
            else:
                return None, array.shape
        tensor_device, tensor_shape = shape_and_device(tensor)
        self_device, self_shape = shape_and_device(self.value)

        device_mismatch_error = f"Can not replicate a variable that is currently on " \
                                f"{self_device} devices to {tensor_device} devices."
        assert (tensor_device is None) or (self_device is None) or (self_device == tensor_device), device_mismatch_error

        tracer_types = (jax.interpreters.partial_eval.JaxprTracer,
                        jax.interpreters.partial_eval.DynamicJaxprTracer)

        shorter_length = min(len(tensor.shape), len(self.value.shape))
        is_special_ok = (isinstance(tensor, tracer_types) or isinstance(self.value, tracer_types))
        is_special_ok = is_special_ok and self.value.shape[-shorter_length:] == tensor.shape[-shorter_length:]

        shape_mismatch_error = f"Assign can not change shape of variable. The current variable shape is {self_shape}," \
                               f" but the requested new shape is {tensor_shape}."
{'fraction_non_alphanumeric': 0.0874158193582464, 'fraction_numerical': 0.0029050574409084907, 'mean_word_length': 3.0676691729323307, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5065601', 'n_tokens_mistral': 2146, 'n_tokens_neox': 2085, 'n_words': 709}