<commit_before># Copyright 2018 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.

import numpy as np
from typing import Any, Callable, Dict, Optional, Tuple, Union

import jax
from ..config import config
from .. import core
from ..core import Trace, Tracer
from ..abstract_arrays import ShapedArray, raise_to_shaped
from ..ad_util import add_jaxvals, add_jaxvals_p, zeros_like_jaxval, zeros_like_p
from .. import linear_util as lu
from ..util import (unzip2, partial, safe_map, wrap_name, split_list,
                    canonicalize_axis)
from . import xla
from . import partial_eval as pe

map = safe_map


def batch(fun: lu.WrappedFun, in_vals, in_dims, out_dim_dests, axis_name):
  # executes a batched version of `fun` following out_dim_dests
  batched_fun = batch_fun(fun, in_dims, out_dim_dests, axis_name=axis_name)
  return batched_fun.call_wrapped(*in_vals)

@lu.transformation_with_aux
def batch_subtrace(main, in_dims, *in_vals, **params):
  trace = BatchTrace(main, core.cur_sublevel())
  in_tracers = [BatchTracer(trace, val, dim) if dim is not None else val
                for val, dim in zip(in_vals, in_dims)]
  outs = yield in_tracers, params
  out_tracers = map(trace.full_raise, outs)
  out_vals, out_dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
  yield out_vals, out_dims


def batch_fun(fun : lu.WrappedFun, in_dims, out_dim_dests, axis_name,
              sum_match=False):
  # transformation version of batch, which doesn't call the function
  fun, out_dims = batch_subtrace(fun)
  return _batch_fun(fun, axis_name, sum_match, in_dims, out_dims, out_dim_dests)

@lu.transformation
def _batch_fun(axis_name, sum_match, in_dims, out_dims_thunk, out_dim_dests,
               *in_vals, **params):
  in_dims = in_dims() if callable(in_dims) else in_dims
  in_dims = [
    canonicalize_axis(dim, np.ndim(val)) if isinstance(dim, int) else dim
    for val, dim in zip(in_vals, in_dims)]
  size, = {x.shape[d] for x, d in zip(in_vals, in_dims) if d is not not_mapped}
  with core.new_main(BatchTrace) as main:
    with core.extend_axis_env(axis_name, size, main):
      out_vals = yield (main, in_dims,) + in_vals, params
    del main
  out_dim_dests = out_dim_dests() if callable(out_dim_dests) else out_dim_dests
  out_dims = out_dims_thunk()
  for od, od_dest in zip(out_dims, out_dim_dests):
    if od is not None and not isinstance(od_dest, int) and not od_dest is last and not sum_match:
      msg = f"vmap has mapped output but out_axes is {od_dest}"
      raise ValueError(msg)
  out_vals = map(partial(matchaxis, size, sum_match=sum_match), out_dims, out_dim_dests, out_vals)
  yield out_vals

def batch_fun2(fun : lu.WrappedFun, in_dims):
  # like `batch_fun` but returns output batch dims (so no out_dim_dests)
  fun, out_dims = batch_subtrace(fun)
  return _batch_fun2(fun, in_dims), out_dims

@lu.transformation
def _batch_fun2(in_dims, *in_vals, **params):
  with core.new_main(BatchTrace) as main:
    out_vals = yield (main, in_dims,) + in_vals, params
    del main
  yield out_vals


### tracer

# TODO(mattjj): use a special sentinel type rather than None
NotMapped = type(None)
not_mapped = None

class BatchTracer(Tracer):
  __slots__ = ['val', 'batch_dim']

  def __init__(self, trace, val, batch_dim: Optional[int]):
    assert core.skip_checks or type(batch_dim) in (int, NotMapped)  # type: ignore
    self._trace = trace
    self.val = val
    self.batch_dim = batch_dim

  @property
  def aval(self):
    aval = raise_to_shaped(core.get_aval(self.val))
    if self.batch_dim is not_mapped:
      return aval
    else:
      if aval is core.abstract_unit:
        return aval
      elif type(aval) is ShapedArray:
        assert 0 <= self.batch_dim < aval.ndim
        new_shape = tuple(np.delete(aval.shape, self.batch_dim))
        return ShapedArray(new_shape, aval.dtype)
      else:
        raise TypeError(aval)

  def full_lower(self):
    if self.batch_dim is not_mapped:
      return core.full_lower(self.val)
    else:
      return self

class BatchTrace(Trace):
  def pure(self, val):
    return BatchTracer(self, val, not_mapped)

  def lift(self, val):
    return BatchTracer(self, val, not_mapped)

  def sublift(self, val):
    return BatchTracer(self, val.val, val.batch_dim)

  def process_primitive(self, primitive, tracers, params):
    vals_in, dims_in = unzip2((t.val, t.batch_dim) for t in tracers)
    if all(bdim is not_mapped for bdim in dims_in):
      return primitive.bind(*vals_in, **params)
    elif config.omnistaging_enabled and primitive in collective_rules:
      axis_names = params['axis_name']
      if not isinstance(axis_names, (tuple, list)):
        axis_names = (axis_names,)
      for i, axis_name in enumerate(axis_names):
        frame = core.axis_frame(axis_name)
        if frame.tag is not self.main:
          continue
        # We run the split_axis rule with tracers, which is supposed to never
        # mix this axis name with another one. We will handle any invocations
        # of collectives over the vmapped axis in a recursive call from a tracer.
        if len(axis_names) > 1:
          return split_axis(primitive, axis_name, tracers, params)
        vals_out, dims_out = collective_rules[primitive](vals_in, dims_in, frame.size, **params)
        results = map(partial(BatchTracer, self), vals_out, dims_out)
        print(results)
        return results if primitive.multiple_results else results[0]
    # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    batched_primitive = get_primitive_batcher(primitive)
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    if primitive.multiple_results:
      return map(partial(BatchTracer, self), val_out, dim_out)
    else:
      return BatchTracer(self, val_out, dim_out)

  def process_call(self, call_primitive, f: lu.WrappedFun, tracers, params):
    assert call_primitive.multiple_results
    params = dict(params, name=wrap_name(params.get('name', f.__name__), 'vmap'))
    vals, dims = unzip2((t.val, t.batch_dim) for t in tracers)
    if all(bdim is not_mapped for bdim in dims):
      return call_primitive.bind(f, *vals, **params)
    else:
      f, dims_out = batch_subtrace(f, self.main, dims)
      vals_out = call_primitive.bind(f, *vals, **params)
      return [BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out())]

  def post_process_call(self, call_primitive, out_tracers, params):
    vals, dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
    main = self.main<commit_msg>Fix axis checking and remove extra print statement (#4184)<commit_after># Copyright 2018 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.

import numpy as np
from typing import Any, Callable, Dict, Optional, Tuple, Union

import jax
from ..config import config
from .. import core
from ..core import Trace, Tracer
from ..abstract_arrays import ShapedArray, raise_to_shaped
from ..ad_util import add_jaxvals, add_jaxvals_p, zeros_like_jaxval, zeros_like_p
from .. import linear_util as lu
from ..util import (unzip2, partial, safe_map, wrap_name, split_list,
                    canonicalize_axis)
from . import xla
from . import partial_eval as pe

map = safe_map


def batch(fun: lu.WrappedFun, in_vals, in_dims, out_dim_dests, axis_name):
  # executes a batched version of `fun` following out_dim_dests
  batched_fun = batch_fun(fun, in_dims, out_dim_dests, axis_name=axis_name)
  return batched_fun.call_wrapped(*in_vals)

@lu.transformation_with_aux
def batch_subtrace(main, in_dims, *in_vals, **params):
  trace = BatchTrace(main, core.cur_sublevel())
  in_tracers = [BatchTracer(trace, val, dim) if dim is not None else val
                for val, dim in zip(in_vals, in_dims)]
  outs = yield in_tracers, params
  out_tracers = map(trace.full_raise, outs)
  out_vals, out_dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
  yield out_vals, out_dims


def batch_fun(fun : lu.WrappedFun, in_dims, out_dim_dests, axis_name,
              sum_match=False):
  # transformation version of batch, which doesn't call the function
  fun, out_dims = batch_subtrace(fun)
  return _batch_fun(fun, axis_name, sum_match, in_dims, out_dims, out_dim_dests)

@lu.transformation
def _batch_fun(axis_name, sum_match, in_dims, out_dims_thunk, out_dim_dests,
               *in_vals, **params):
  in_dims = in_dims() if callable(in_dims) else in_dims
  in_dims = [
    canonicalize_axis(dim, np.ndim(val)) if isinstance(dim, int) else dim
    for val, dim in zip(in_vals, in_dims)]
  size, = {x.shape[d] for x, d in zip(in_vals, in_dims) if d is not not_mapped}
  with core.new_main(BatchTrace) as main:
    with core.extend_axis_env(axis_name, size, main):
      out_vals = yield (main, in_dims,) + in_vals, params
    del main
  out_dim_dests = out_dim_dests() if callable(out_dim_dests) else out_dim_dests
  out_dims = out_dims_thunk()
  for od, od_dest in zip(out_dims, out_dim_dests):
    if od is not None and not isinstance(od_dest, int) and not od_dest is last and not sum_match:
      msg = f"vmap has mapped output but out_axes is {od_dest}"
      raise ValueError(msg)
  out_vals = map(partial(matchaxis, size, sum_match=sum_match), out_dims, out_dim_dests, out_vals)
  yield out_vals

def batch_fun2(fun : lu.WrappedFun, in_dims):
  # like `batch_fun` but returns output batch dims (so no out_dim_dests)
  fun, out_dims = batch_subtrace(fun)
  return _batch_fun2(fun, in_dims), out_dims

@lu.transformation
def _batch_fun2(in_dims, *in_vals, **params):
  with core.new_main(BatchTrace) as main:
    out_vals = yield (main, in_dims,) + in_vals, params
    del main
  yield out_vals


### tracer

# TODO(mattjj): use a special sentinel type rather than None
NotMapped = type(None)
not_mapped = None

class BatchTracer(Tracer):
  __slots__ = ['val', 'batch_dim']

  def __init__(self, trace, val, batch_dim: Optional[int]):
    assert core.skip_checks or type(batch_dim) in (int, NotMapped)  # type: ignore
    self._trace = trace
    self.val = val
    self.batch_dim = batch_dim

  @property
  def aval(self):
    aval = raise_to_shaped(core.get_aval(self.val))
    if self.batch_dim is not_mapped:
      return aval
    else:
      if aval is core.abstract_unit:
        return aval
      elif type(aval) is ShapedArray:
        assert 0 <= self.batch_dim < aval.ndim
        new_shape = tuple(np.delete(aval.shape, self.batch_dim))
        return ShapedArray(new_shape, aval.dtype)
      else:
        raise TypeError(aval)

  def full_lower(self):
    if self.batch_dim is not_mapped:
      return core.full_lower(self.val)
    else:
      return self

class BatchTrace(Trace):
  def pure(self, val):
    return BatchTracer(self, val, not_mapped)

  def lift(self, val):
    return BatchTracer(self, val, not_mapped)

  def sublift(self, val):
    return BatchTracer(self, val.val, val.batch_dim)

  def process_primitive(self, primitive, tracers, params):
    vals_in, dims_in = unzip2((t.val, t.batch_dim) for t in tracers)
    if all(bdim is not_mapped for bdim in dims_in):
      return primitive.bind(*vals_in, **params)
    elif config.omnistaging_enabled and primitive in collective_rules:
      axis_names = params['axis_name']
      if not isinstance(axis_names, (tuple, list)):
        axis_names = (axis_names,)
      for i, axis_name in enumerate(axis_names):
        frame = core.axis_frame(axis_name)
        if frame.main_trace is not self.main:
          continue
        # We run the split_axis rule with tracers, which is supposed to never
        # mix this axis name with another one. We will handle any invocations
        # of collectives over the vmapped axis in a recursive call from a tracer.
        if len(axis_names) > 1:
          return split_axis(primitive, axis_name, tracers, params)
        vals_out, dims_out = collective_rules[primitive](vals_in, dims_in, frame.size, **params)
        results = map(partial(BatchTracer, self), vals_out, dims_out)
        return results if primitive.multiple_results else results[0]
    # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    batched_primitive = get_primitive_batcher(primitive)
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    if primitive.multiple_results:
      return map(partial(BatchTracer, self), val_out, dim_out)
    else:
      return BatchTracer(self, val_out, dim_out)

  def process_call(self, call_primitive, f: lu.WrappedFun, tracers, params):
    assert call_primitive.multiple_results
    params = dict(params, name=wrap_name(params.get('name', f.__name__), 'vmap'))
    vals, dims = unzip2((t.val, t.batch_dim) for t in tracers)
    if all(bdim is not_mapped for bdim in dims):
      return call_primitive.bind(f, *vals, **params)
    else:
      f, dims_out = batch_subtrace(f, self.main, dims)
      vals_out = call_primitive.bind(f, *vals, **params)
      return [BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out())]

  def post_process_call(self, call_primitive, out_tracers, params):
    vals, dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
    main = self.main
{'fraction_non_alphanumeric': 0.09229129662522202, 'fraction_numerical': 0.002984014209591474, 'mean_word_length': 3.9045296167247385, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 2, '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': '6618131', 'n_tokens_mistral': 4849, 'n_tokens_neox': 4798, 'n_words': 1420}