<commit_before># coding=utf-8
# Copyright 2021 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
#
#     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.

# Lint as: python3
"""Implements a geometry class for points supported on a cartesian product."""
import itertools
from typing import Optional, Sequence, Union

import jax
import jax.numpy as jnp
import numpy as np
from ott.geometry import costs
from ott.geometry import geometry
from ott.geometry import ops


@jax.tree_util.register_pytree_node_class
class Grid(geometry.Geometry):
  r"""Class describing the geometry of points taken in a cartestian product.

  This class implements a geometry in which probability measures are supported
  on a :math:`d`-dimensional cartesian grid, a cartesian product of :math:`d`
  lists of values, each list being itself of size :math:`n_i`.

  The transportation cost between points in the grid is assumed to be separable,
  namely a sum of coordinate-wise cost functions, as in
    :math:`cost(x,y) = \sum_{i=1}^d cost_i(x_i, y_i)` where :math:`cost_i`: R x R → R.

  In such a regime, and despite the fact that the total number :math:`n_{total}`
  of points in the grid is exponential :math:`d` (namely :math:`\prod_i n_i`),
  applying a kernel in the context of regularized optimal transport can be
  carried out in time that is of the order of :math:`n_{total}^{(1+1/d)}` using
  convolutions, either in the original domain or log-space domain. This class
  precomputes :math:`d` :math:`n_i` x :math:`n_i` cost matrices (one per
  dimension) and implements these two operations by carrying out these
  convolutions one dimension at a time.
  """

  def __init__(
      self,
      x: Optional[Sequence[jnp.ndarray]] = None,
      grid_size: Optional[Sequence[int]] = None,
      cost_fns: Optional[Sequence[costs.CostFn]] = None,
      num_a: Optional[int] = None,
      grid_dimension: int = None,
      **kwargs):
    """Create instance of grid using either locations or sizes.

    Args:
      x : list of arrays of varying sizes, describing the locations of the grid.
        Locations are provided as a list of jnp.ndarrays, that is :math:`d`
        vectors of (possibly varying) size :math:`n_i`. The resulting grid
        is the Cartesian product of these vectors.
      grid_size: t-uple of integers describing grid sizes, namely
        :math:`(n_1,...,n_d)`. This will only be used if x is None.
        In that case the grid will be assumed to lie in the hypercube
        :math:`[0,1]^d`, with the :math:`d` dimensions, described as points
        regularly sampled in [0,1].
      cost_fns: a sequence of :math:`d` costs.CostFn's, each being a cost taking
        two reals as inputs to output a real number.
      num_a: total size of grid. This parameters will be computed from other
        inputs and used in the flatten/unflatten functions.
      grid_dimension: dimension of grid. This parameters will be computed from
        other inputs and used in the flatten/unflatten functions.
      **kwargs: other optional parameters to be passed on to superclass
        initializer, notably those related to epsilon regularization.

    """
    if (grid_size is not None and x is not None and num_a is not None and
        grid_dimension is not None):
      self.grid_size = grid_size
      self.x = x
      self.num_a = num_a
      self.grid_dimension = grid_dimension
    elif x is not None:
      self.x = x
      self.grid_size = tuple([xs.shape[0] for xs in x])
      self.num_a = np.prod(np.array(self.grid_size))
      self.grid_dimension = len(self.x)
    elif grid_size is not None:
      self.grid_size = grid_size
      self.x = tuple([jnp.linspace(0, 1, n) for n in self.grid_size])
      self.num_a = np.prod(np.array(grid_size))
      self.grid_dimension = len(self.grid_size)
    else:
      raise ValueError('Input either grid_size t-uple or grid locations x.')

    if cost_fns is None:
      cost_fns = [costs.Euclidean()]
    self.cost_fns = cost_fns<commit_msg>[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.<commit_after># coding=utf-8
# Copyright 2021 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
#
#     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.

# Lint as: python3
"""Implements a geometry class for points supported on a cartesian product."""
import itertools
from typing import Optional, Sequence, Union

import jax
import jax.numpy as jnp
import numpy as np
from ott.geometry import costs
from ott.geometry import geometry
from ott.geometry import ops


@jax.tree_util.register_pytree_node_class
class Grid(geometry.Geometry):
  r"""Class describing the geometry of points taken in a cartestian product.

  This class implements a geometry in which probability measures are supported
  on a :math:`d`-dimensional cartesian grid, a cartesian product of :math:`d`
  lists of values, each list being itself of size :math:`n_i`.

  The transportation cost between points in the grid is assumed to be separable,
  namely a sum of coordinate-wise cost functions, as in
    :math:`cost(x,y) = \sum_{i=1}^d cost_i(x_i, y_i)` where :math:`cost_i`: R x R → R.

  In such a regime, and despite the fact that the total number :math:`n_{total}`
  of points in the grid is exponential :math:`d` (namely :math:`\prod_i n_i`),
  applying a kernel in the context of regularized optimal transport can be
  carried out in time that is of the order of :math:`n_{total}^{(1+1/d)}` using
  convolutions, either in the original domain or log-space domain. This class
  precomputes :math:`d` :math:`n_i` x :math:`n_i` cost matrices (one per
  dimension) and implements these two operations by carrying out these
  convolutions one dimension at a time.
  """

  def __init__(
      self,
      x: Optional[Sequence[jnp.ndarray]] = None,
      grid_size: Optional[Sequence[int]] = None,
      cost_fns: Optional[Sequence[costs.CostFn]] = None,
      num_a: Optional[int] = None,
      grid_dimension: int = None,
      **kwargs):
    """Create instance of grid using either locations or sizes.

    Args:
      x : list of arrays of varying sizes, describing the locations of the grid.
        Locations are provided as a list of jnp.ndarrays, that is :math:`d`
        vectors of (possibly varying) size :math:`n_i`. The resulting grid
        is the Cartesian product of these vectors.
      grid_size: t-uple of integers describing grid sizes, namely
        :math:`(n_1,...,n_d)`. This will only be used if x is None.
        In that case the grid will be assumed to lie in the hypercube
        :math:`[0,1]^d`, with the :math:`d` dimensions, described as points
        regularly sampled in [0,1].
      cost_fns: a sequence of :math:`d` costs.CostFn's, each being a cost taking
        two reals as inputs to output a real number.
      num_a: total size of grid. This parameters will be computed from other
        inputs and used in the flatten/unflatten functions.
      grid_dimension: dimension of grid. This parameters will be computed from
        other inputs and used in the flatten/unflatten functions.
      **kwargs: other optional parameters to be passed on to superclass
        initializer, notably those related to epsilon regularization.

    """
    if (grid_size is not None and x is not None and num_a is not None and
        grid_dimension is not None):
      self.grid_size = tuple(map(int, grid_size))
      self.x = x
      self.num_a = num_a
      self.grid_dimension = grid_dimension
    elif x is not None:
      self.x = x
      self.grid_size = tuple([xs.shape[0] for xs in x])
      self.num_a = np.prod(np.array(self.grid_size))
      self.grid_dimension = len(self.x)
    elif grid_size is not None:
      self.grid_size = tuple(map(int, grid_size))
      self.x = tuple([jnp.linspace(0, 1, n) for n in self.grid_size])
      self.num_a = np.prod(np.array(grid_size))
      self.grid_dimension = len(self.grid_size)
    else:
      raise ValueError('Input either grid_size t-uple or grid locations x.')

    if cost_fns is None:
      cost_fns = [costs.Euclidean()]
    self.cost_fns = cost_fns
{'fraction_non_alphanumeric': 0.08199801082992596, 'fraction_numerical': 0.004641396839429771, 'mean_word_length': 3.376208897485493, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '6774003', 'n_tokens_mistral': 2659, 'n_tokens_neox': 2609, 'n_words': 1139}