<commit_before># Copyright 2022 The jax_triton Authors.
#
# 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.

"""Module for pallas-specific JAX primitives and functions."""
from __future__ import annotations
import enum
import functools

from typing import Any, List, Optional, Tuple

import jax
from jax import core as jax_core
from jax import lax
from jax._src import ad_util
from jax._src.util import (safe_map, safe_zip, split_list, merge_lists,
                           partition_list)
from jax._src.state import primitives as state_primitives
from jax.interpreters import ad
import jax.numpy as jnp
import numpy as np

map, unsafe_map = safe_map, map
zip, unsafe_zip = safe_zip, zip

def _process_idx(idx, ref_shape):
  if any(isinstance(i, slice) and i != slice(None) for i in idx):
    raise NotImplementedError("Non-`slice(None)` slices not supported yet.")
  if len(idx) != len(ref_shape):
    raise ValueError("Must provide indexer for each dimension of `Ref`.")
  is_int_indexing = [isinstance(i, (jnp.ndarray, int)) for i in idx]
  other_indexers, int_indexers = partition_list(is_int_indexing, idx)
  int_indexers = [np.array(i, np.int32) if isinstance(i, int) else i for i in
                  int_indexers]
  indexer_shapes = [jnp.shape(i) for i in int_indexers]
  bcast_shape = tuple(s for i in indexer_shapes for s in i)
  idx_iter = iter(range(len(bcast_shape)))
  int_indexers = [
      lax.broadcast_in_dim(i, bcast_shape, tuple(next(idx_iter) for _ in
                                                 range(len(i.shape))))
      for i in int_indexers
  ]
  return merge_lists(is_int_indexing, other_indexers, int_indexers)

program_id_p = jax_core.Primitive("program_id")

def program_id(axis):
  return program_id_p.bind(axis=axis)

def _program_id_abstract_eval(**_):
  return jax_core.ShapedArray((), jnp.int32)
program_id_p.def_abstract_eval(_program_id_abstract_eval)

class AtomicOpType(enum.Enum):
  XCHG = "xchg"
  ADD = "add"
  MAX = "max"
  MIN = "min"
  AND = "and"
  OR = "or"
  XOR = "xor"

atomic_rmw_p = jax_core.Primitive("atomic_rmw")

def _atomic_abstract_eval(ref_aval, val_aval, *all_avals,
                          indexed_dims: Tuple[bool], atomic_type: AtomicOpType,
                          **_: Any):
  if ref_aval.dtype == jnp.dtype("float16") and atomic_type != AtomicOpType.ADD:
    raise ValueError(f"`atomic_{atomic_type.value}` does not support f16.")
  if ref_aval.dtype in {jnp.dtype("bool"), jnp.dtype("int8"),
                        jnp.dtype("int16"), jnp.bfloat16}:
    raise ValueError(f"`atomic_{atomic_type.value}` does not support {ref_aval.dtype}.")
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._swap_abstract_eval(
      ref_aval, val_aval, *idx_avals, indexed_dims=indexed_dims)
atomic_rmw_p.def_effectful_abstract_eval(_atomic_abstract_eval)

def atomic_rmw(x_ref, idx, val, *, mask: Optional[Any] = None, atomic_type: AtomicOpType):
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:
    args = (*args, mask)
  return atomic_rmw_p.bind(x_ref, val, *args, indexed_dims=indexed_dims,
                           atomic_type=atomic_type, masked=mask is not None)

atomic_xchg = functools.partial(atomic_rmw, atomic_type=AtomicOpType.XCHG)
atomic_add = functools.partial(atomic_rmw, atomic_type=AtomicOpType.ADD)
atomic_max = functools.partial(atomic_rmw, atomic_type=AtomicOpType.MAX)
atomic_min = functools.partial(atomic_rmw, atomic_type=AtomicOpType.MIN)
atomic_and = functools.partial(atomic_rmw, atomic_type=AtomicOpType.AND)
atomic_or = functools.partial(atomic_rmw, atomic_type=AtomicOpType.OR)
atomic_xor = functools.partial(atomic_rmw, atomic_type=AtomicOpType.XOR)

max_contiguous_p = jax_core.Primitive("max_contiguous")

def max_contiguous(x, values):
  if not isinstance(values, list):
    values = [values]
  return max_contiguous_p.bind(x, values=values)

def _max_contiguous_abstract_eval(aval, **_):
  return aval
max_contiguous_p.def_abstract_eval(_max_contiguous_abstract_eval)

multiple_of_p = jax_core.Primitive("multiple_of")

def multiple_of(x, values):
  if not isinstance(values, list):
    values = [values]
  return multiple_of_p.bind(x, values=values)

def _multiple_of_abstract_eval(aval, **_):
  return aval
multiple_of_p.def_abstract_eval(_multiple_of_abstract_eval)

load_p = jax_core.Primitive('masked_load')

def _load_abstract_eval(ref_aval, *all_avals, indexed_dims: Tuple[bool],
                        **params: Any):
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._get_abstract_eval(
      ref_aval, *idx_avals, indexed_dims=indexed_dims)
load_p.def_effectful_abstract_eval(_load_abstract_eval)

def _load_jvp(primals: List[Any], tangents: List[Any], *, indexed_dims, **params):
  ref_primal, *idx_and_rest_primals = primals
  ref_tangent, *idx_and_rest_tangents = tangents
  idx_primals, _ = split_list(idx_and_rest_primals, [sum(indexed_dims)])
  _, rest_tangents = split_list(idx_and_rest_tangents, [sum(indexed_dims)])
  return (load_p.bind(ref_primal, *idx_and_rest_primals,
                      indexed_dims=indexed_dims, **params),
          load_p.bind(ref_tangent, *idx_primals, *rest_tangents,
                      indexed_dims=indexed_dims, **params))
ad.primitive_jvps[load_p] = _load_jvp

def _load_transpose(g, ref, *idx_and_rest, **params):
  raise NotImplementedError
ad.primitive_transposes[load_p] = _load_transpose

swap_p = jax_core.Primitive('masked_swap')

def _swap_abstract_eval(ref_aval, val_aval, *all_avals, indexed_dims: Tuple[bool],
                        **_: Any):
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._swap_abstract_eval(
      ref_aval, val_aval, *idx_avals, indexed_dims=indexed_dims)
swap_p.def_effectful_abstract_eval(_swap_abstract_eval)

def _swap_jvp(primals: List[Any], tangents: List[Any], *, indexed_dims, **params):
  ref_primal, val_primal, *idx_and_rest_primals = primals
  ref_tangent, val_tangent, *idx_and_rest_tangents = tangents
  idx_primals, _ = split_list(idx_and_rest_primals, [sum(indexed_dims)])
  _, rest_tangents = split_list(idx_and_rest_tangents, [sum(indexed_dims)])
  val_tangent = ad_util.instantiate(val_tangent)
  return (swap_p.bind(ref_primal, val_primal, *idx_and_rest_primals,
                      indexed_dims=indexed_dims, **params),
          swap_p.bind(ref_tangent, val_tangent, *idx_primals, *rest_tangents,
                      indexed_dims=indexed_dims, **params))
ad.primitive_jvps[swap_p] = _swap_jvp

def _swap_transpose(g, ref, *idx_and_rest, **params):
  raise NotImplementedError
ad.primitive_transposes[swap_p] = _swap_transpose

def load(x_ref, idx, *, mask=None, other=None, cache_modifier="",
         eviction_poicy="", volatile=False):
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:
    args = (*args, mask)
  if other is not None:
    assert mask is not None
    args = (*args, other)
  return load_p.bind(x_ref, *args, masked=mask is not None, cache_modifier=cache_modifier,
                     eviction_policy=eviction_poicy, is_volatile=volatile,
                     indexed_dims=indexed_dims)


def swap(x_ref, idx, val, *, mask=None, eviction_policy="") -> Any:
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:<commit_msg>Fix mispelling eviction_poicy -> eviction_policy<commit_after># Copyright 2022 The jax_triton Authors.
#
# 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.

"""Module for pallas-specific JAX primitives and functions."""
from __future__ import annotations
import enum
import functools

from typing import Any, List, Optional, Tuple

import jax
from jax import core as jax_core
from jax import lax
from jax._src import ad_util
from jax._src.util import (safe_map, safe_zip, split_list, merge_lists,
                           partition_list)
from jax._src.state import primitives as state_primitives
from jax.interpreters import ad
import jax.numpy as jnp
import numpy as np

map, unsafe_map = safe_map, map
zip, unsafe_zip = safe_zip, zip

def _process_idx(idx, ref_shape):
  if any(isinstance(i, slice) and i != slice(None) for i in idx):
    raise NotImplementedError("Non-`slice(None)` slices not supported yet.")
  if len(idx) != len(ref_shape):
    raise ValueError("Must provide indexer for each dimension of `Ref`.")
  is_int_indexing = [isinstance(i, (jnp.ndarray, int)) for i in idx]
  other_indexers, int_indexers = partition_list(is_int_indexing, idx)
  int_indexers = [np.array(i, np.int32) if isinstance(i, int) else i for i in
                  int_indexers]
  indexer_shapes = [jnp.shape(i) for i in int_indexers]
  bcast_shape = tuple(s for i in indexer_shapes for s in i)
  idx_iter = iter(range(len(bcast_shape)))
  int_indexers = [
      lax.broadcast_in_dim(i, bcast_shape, tuple(next(idx_iter) for _ in
                                                 range(len(i.shape))))
      for i in int_indexers
  ]
  return merge_lists(is_int_indexing, other_indexers, int_indexers)

program_id_p = jax_core.Primitive("program_id")

def program_id(axis):
  return program_id_p.bind(axis=axis)

def _program_id_abstract_eval(**_):
  return jax_core.ShapedArray((), jnp.int32)
program_id_p.def_abstract_eval(_program_id_abstract_eval)

class AtomicOpType(enum.Enum):
  XCHG = "xchg"
  ADD = "add"
  MAX = "max"
  MIN = "min"
  AND = "and"
  OR = "or"
  XOR = "xor"

atomic_rmw_p = jax_core.Primitive("atomic_rmw")

def _atomic_abstract_eval(ref_aval, val_aval, *all_avals,
                          indexed_dims: Tuple[bool], atomic_type: AtomicOpType,
                          **_: Any):
  if ref_aval.dtype == jnp.dtype("float16") and atomic_type != AtomicOpType.ADD:
    raise ValueError(f"`atomic_{atomic_type.value}` does not support f16.")
  if ref_aval.dtype in {jnp.dtype("bool"), jnp.dtype("int8"),
                        jnp.dtype("int16"), jnp.bfloat16}:
    raise ValueError(f"`atomic_{atomic_type.value}` does not support {ref_aval.dtype}.")
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._swap_abstract_eval(
      ref_aval, val_aval, *idx_avals, indexed_dims=indexed_dims)
atomic_rmw_p.def_effectful_abstract_eval(_atomic_abstract_eval)

def atomic_rmw(x_ref, idx, val, *, mask: Optional[Any] = None, atomic_type: AtomicOpType):
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:
    args = (*args, mask)
  return atomic_rmw_p.bind(x_ref, val, *args, indexed_dims=indexed_dims,
                           atomic_type=atomic_type, masked=mask is not None)

atomic_xchg = functools.partial(atomic_rmw, atomic_type=AtomicOpType.XCHG)
atomic_add = functools.partial(atomic_rmw, atomic_type=AtomicOpType.ADD)
atomic_max = functools.partial(atomic_rmw, atomic_type=AtomicOpType.MAX)
atomic_min = functools.partial(atomic_rmw, atomic_type=AtomicOpType.MIN)
atomic_and = functools.partial(atomic_rmw, atomic_type=AtomicOpType.AND)
atomic_or = functools.partial(atomic_rmw, atomic_type=AtomicOpType.OR)
atomic_xor = functools.partial(atomic_rmw, atomic_type=AtomicOpType.XOR)

max_contiguous_p = jax_core.Primitive("max_contiguous")

def max_contiguous(x, values):
  if not isinstance(values, list):
    values = [values]
  return max_contiguous_p.bind(x, values=values)

def _max_contiguous_abstract_eval(aval, **_):
  return aval
max_contiguous_p.def_abstract_eval(_max_contiguous_abstract_eval)

multiple_of_p = jax_core.Primitive("multiple_of")

def multiple_of(x, values):
  if not isinstance(values, list):
    values = [values]
  return multiple_of_p.bind(x, values=values)

def _multiple_of_abstract_eval(aval, **_):
  return aval
multiple_of_p.def_abstract_eval(_multiple_of_abstract_eval)

load_p = jax_core.Primitive('masked_load')

def _load_abstract_eval(ref_aval, *all_avals, indexed_dims: Tuple[bool],
                        **params: Any):
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._get_abstract_eval(
      ref_aval, *idx_avals, indexed_dims=indexed_dims)
load_p.def_effectful_abstract_eval(_load_abstract_eval)

def _load_jvp(primals: List[Any], tangents: List[Any], *, indexed_dims, **params):
  ref_primal, *idx_and_rest_primals = primals
  ref_tangent, *idx_and_rest_tangents = tangents
  idx_primals, _ = split_list(idx_and_rest_primals, [sum(indexed_dims)])
  _, rest_tangents = split_list(idx_and_rest_tangents, [sum(indexed_dims)])
  return (load_p.bind(ref_primal, *idx_and_rest_primals,
                      indexed_dims=indexed_dims, **params),
          load_p.bind(ref_tangent, *idx_primals, *rest_tangents,
                      indexed_dims=indexed_dims, **params))
ad.primitive_jvps[load_p] = _load_jvp

def _load_transpose(g, ref, *idx_and_rest, **params):
  raise NotImplementedError
ad.primitive_transposes[load_p] = _load_transpose

swap_p = jax_core.Primitive('masked_swap')

def _swap_abstract_eval(ref_aval, val_aval, *all_avals, indexed_dims: Tuple[bool],
                        **_: Any):
  idx_avals, _ = split_list(all_avals, [sum(indexed_dims)])
  return state_primitives._swap_abstract_eval(
      ref_aval, val_aval, *idx_avals, indexed_dims=indexed_dims)
swap_p.def_effectful_abstract_eval(_swap_abstract_eval)

def _swap_jvp(primals: List[Any], tangents: List[Any], *, indexed_dims, **params):
  ref_primal, val_primal, *idx_and_rest_primals = primals
  ref_tangent, val_tangent, *idx_and_rest_tangents = tangents
  idx_primals, _ = split_list(idx_and_rest_primals, [sum(indexed_dims)])
  _, rest_tangents = split_list(idx_and_rest_tangents, [sum(indexed_dims)])
  val_tangent = ad_util.instantiate(val_tangent)
  return (swap_p.bind(ref_primal, val_primal, *idx_and_rest_primals,
                      indexed_dims=indexed_dims, **params),
          swap_p.bind(ref_tangent, val_tangent, *idx_primals, *rest_tangents,
                      indexed_dims=indexed_dims, **params))
ad.primitive_jvps[swap_p] = _swap_jvp

def _swap_transpose(g, ref, *idx_and_rest, **params):
  raise NotImplementedError
ad.primitive_transposes[swap_p] = _swap_transpose

def load(x_ref, idx, *, mask=None, other=None, cache_modifier="",
         eviction_policy="", volatile=False):
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:
    args = (*args, mask)
  if other is not None:
    assert mask is not None
    args = (*args, other)
  return load_p.bind(x_ref, *args, masked=mask is not None, cache_modifier=cache_modifier,
                     eviction_policy=eviction_policy, is_volatile=volatile,
                     indexed_dims=indexed_dims)


def swap(x_ref, idx, val, *, mask=None, eviction_policy="") -> Any:
  idx = _process_idx(idx, x_ref.shape)
  idx, indexed_dims = state_primitives._unpack_idx(idx, x_ref.ndim)
  args = idx
  if mask is not None:
{'fraction_non_alphanumeric': 0.10193059780247067, 'fraction_numerical': 0.002607238189831771, 'mean_word_length': 4.60933147632312, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4988020', 'n_tokens_mistral': 5790, 'n_tokens_neox': 5777, 'n_words': 1214}